Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1from builtins import zip 

2from builtins import map 

3from builtins import range 

4# nd Slicer slices data on N columns in simData 

5 

6import warnings 

7import numpy as np 

8import matplotlib.pyplot as plt 

9from matplotlib import colors 

10import itertools 

11from functools import wraps 

12 

13from lsst.sims.maf.plots.ndPlotters import TwoDSubsetData, OneDSubsetData 

14from .baseSlicer import BaseSlicer 

15 

16__all__ = ['NDSlicer'] 

17 

18class NDSlicer(BaseSlicer): 

19 """Nd slicer (N dimensions)""" 

20 def __init__(self, sliceColList=None, verbose=True, binsList=100): 

21 """Instantiate object. 

22 binsList can be a list of numpy arrays with the respective slicepoints for sliceColList, 

23 or a list of integers (one per column in sliceColList) or a single value 

24 (repeated for all columns, default=100).""" 

25 super(NDSlicer, self).__init__(verbose=verbose) 

26 self.bins = None 

27 self.nslice = None 

28 self.sliceColList = sliceColList 

29 self.columnsNeeded = self.sliceColList 

30 if self.sliceColList is not None: 

31 self.nD = len(self.sliceColList) 

32 else: 

33 self.nD = None 

34 self.binsList = binsList 

35 if not (isinstance(binsList, float) or isinstance(binsList, int)): 

36 if len(self.binsList) != self.nD: 

37 raise Exception('BinsList must be same length as sliceColNames, unless it is a single value') 

38 self.slicer_init={'sliceColList':sliceColList} 

39 self.plotFuncs = [TwoDSubsetData, OneDSubsetData] 

40 

41 def setupSlicer(self, simData, maps=None): 

42 """Set up bins. """ 

43 # Parse input bins choices. 

44 self.bins = [] 

45 # If we were given a single number for the binsList, convert to list. 

46 if isinstance(self.binsList, float) or isinstance(self.binsList, int): 

47 self.binsList = [self.binsList for c in self.sliceColList] 

48 # And then build bins. 

49 for bl, col in zip(self.binsList, self.sliceColList): 

50 if isinstance(bl, float) or isinstance(bl, int): 

51 sliceCol = simData[col] 

52 binMin = sliceCol.min() 

53 binMax = sliceCol.max() 

54 if binMin == binMax: 

55 warnings.warn('BinMin=BinMax for column %s: increasing binMax by 1.' %(col)) 

56 binMax = binMax + 1 

57 binsize = (binMax - binMin) / float(bl) 

58 bins = np.arange(binMin, binMax + binsize/2.0, binsize, 'float') 

59 self.bins.append(bins) 

60 else: 

61 self.bins.append(np.sort(bl)) 

62 # Count how many bins we have total (not counting last 'RHS' bin values, as in oneDSlicer). 

63 self.nslice = (np.array(list(map(len, self.bins)))-1).prod() 

64 # Set up slice metadata. 

65 self.slicePoints['sid'] = np.arange(self.nslice) 

66 # Including multi-D 'leftmost' bin values 

67 binsForIteration = [] 

68 for b in self.bins: 

69 binsForIteration.append(b[:-1]) 

70 biniterator = itertools.product(*binsForIteration) 

71 self.slicePoints['bins'] = [] 

72 for b in biniterator: 

73 self.slicePoints['bins'].append(b) 

74 # and multi-D 'leftmost' bin indexes corresponding to each sid 

75 self.slicePoints['binIdxs'] = [] 

76 binIdsForIteration = [] 

77 for b in self.bins: 

78 binIdsForIteration.append(np.arange(len(b[:-1]))) 

79 binIdIterator = itertools.product(*binIdsForIteration) 

80 for bidx in binIdIterator: 

81 self.slicePoints['binIdxs'].append(bidx) 

82 # Add metadata from maps. 

83 self._runMaps(maps) 

84 # Set up indexing for data slicing. 

85 self.simIdxs = [] 

86 self.lefts = [] 

87 for sliceColName, bins in zip(self.sliceColList, self.bins): 

88 simIdxs = np.argsort(simData[sliceColName]) 

89 simFieldsSorted = np.sort(simData[sliceColName]) 

90 # "left" values are location where simdata == bin value 

91 left = np.searchsorted(simFieldsSorted, bins[:-1], 'left') 

92 left = np.concatenate((left, np.array([len(simIdxs),]))) 

93 # Add these calculated values into the class lists of simIdxs and lefts. 

94 self.simIdxs.append(simIdxs) 

95 self.lefts.append(left) 

96 @wraps (self._sliceSimData) 

97 def _sliceSimData(islice): 

98 """Slice simData to return relevant indexes for slicepoint.""" 

99 # Identify relevant pointings in each dimension. 

100 simIdxsList = [] 

101 # Translate islice into indexes in each bin dimension 

102 binIdxs = self.slicePoints['binIdxs'][islice] 

103 for d, i in zip(list(range(self.nD)), binIdxs): 

104 simIdxsList.append(set(self.simIdxs[d][self.lefts[d][i]:self.lefts[d][i+1]])) 

105 idxs = list(set.intersection(*simIdxsList)) 

106 return {'idxs':idxs, 

107 'slicePoint':{'sid':islice, 

108 'binLeft':self.slicePoints['bins'][islice], 

109 'binIdx':self.slicePoints['binIdxs'][islice]}} 

110 setattr(self, '_sliceSimData', _sliceSimData) 

111 

112 def __eq__(self, otherSlicer): 

113 """Evaluate if grids are equivalent.""" 

114 if isinstance(otherSlicer, NDSlicer): 

115 if otherSlicer.nD != self.nD: 

116 return False 

117 for i in range(self.nD): 

118 if np.all(otherSlicer.slicePoints['bins'][i] != self.slicePoints['bins'][i]): 

119 return False 

120 return True 

121 else: 

122 return False