Coverage for webgeodyn/data/measuredata.py: 0%
136 statements
« prev ^ index » next coverage.py v7.2.7, created at 2023-12-18 09:33 +0000
1import numpy as np
2from scipy.special import lpmn
5class MeasureData():
6 """
7 Mother class of GHData and TSData. Handles the measure data.
8 """
9 def __array_wrap__(self, result):
10 """
11 Reimplementation of __array_wrap__ Numpy method.
13 :param result: result to add to data
14 :type result: any
15 :return: data + result
16 :rtype: np.array
17 """
18 return np.array(self.data + result)
20 def __init__(self, shape, lmax, PnmNorm,
21 units=None, gridth=None, gridph=None):
22 """
23 Constructor of MeasureData. Set all members.
25 :param shape: shape of the data
26 :type shape: tuple
27 :param lmax: maximum degree of the measure
28 :type lmax: int
29 :param PnmNorm: normalisation to use for Legendre polynomials
30 :type PnmNorm: function
31 :param units: unit of measure (default=None)
32 :type units: str
33 :param gridth: grid where thetas are evaluated (default=None)
34 :type gridth: np.array
35 :param gridph: grid where phis are evaluated (default=None)
36 :type gridph: np.array
37 """
38 self.data = np.zeros(shape)
39 self.lmax = lmax
40 self.PnmNorm = PnmNorm
41 self.computingState = 0
42 self.cachedRThetaPhiData = {}
43 self.cachedRThetaPhiData_mean = {}
44 self.units = units
45 self.has_grid = False
46 self.th = None
47 self.ph = None
48 if gridth is not None and gridph is not None:
49 self.setGrid(gridth, gridph)
51 def truncate(self, lmax):
52 """ Updates the data to keep only the data corresponding to l<lmax
53 :param lmax: new maximum degree
54 :type lmax: int
55 """
56 if lmax > self.lmax:
57 print("WARNING: truncation should be less that %i"%self.lmax)
58 return
59 self.lmax = lmax
60 self.initMeasure()
61 self.setData(self.data)
63 def setGrid(self, gridth, gridph):
64 """ Saves the theta and phi grids (and their dimensions) given when building the data.
65 :param gridth: grid where thetas are evaluated
66 :type gridth: np.array
67 :param gridph: grid where phis are evaluated
68 :type gridph: np.array
69 """
70 self.th = gridth
71 self.ph = gridph
72 self.nth = self.th.shape[0]
73 self.nph = self.ph.shape[0]
74 self.has_grid = True
76 def computePnm(self, th):
77 """ Computes the associated Legendre polynomials for a given theta.
78 :param th: theta where Pnm are evaluated
79 :type th: float
80 :return: list of the Legendre polynomial
81 :rtype:
82 """
83 z = np.cos(th)
84 # Scipy function lpmn computes the Legendre polynomials
85 # and their derivative for all orders 0..lmax and degrees 0..lmax
86 Pnm, dPnm = lpmn(self.lmax, self.lmax, z)
87 # Apply the chosen normalisation on Pnm
88 # TODO : Check if normalisation exists
89 norm_constants = self.PnmNorm(self.lmax)
90 Pnm *= norm_constants
91 # The derivative should be given wrt. theta (python gives it wrt. z)
92 dPnm *= -np.sin(th)*norm_constants
94 # Reshape (l,m) to k
95 Pnmk = Pnm[self.m, self.l]
96 dPnmk = dPnm[self.m, self.l]
98 return Pnmk, dPnmk
100 def computeGridPnm(self):
101 """ Computes associated Legendre polynomials on the theta grid stored in self.th. """
102 if not self.has_grid:
103 raise ValueError("Can not compute Pnm along grid, data has no defined real space (th,ph) grid")
104 self.Pnm = np.zeros((self.nth, self.nk))
105 self.dPnm = np.zeros((self.nth, self.nk))
106 for i_theta, theta in enumerate(self.th):
107 self.Pnm[i_theta], self.dPnm[i_theta] = self.computePnm(theta)
109 def clearCache(self):
110 """ Clears the cached data (usually called when setting new data) """
111 self.cachedRThetaPhiData = {}
112 self.cachedRThetaPhiData_mean = {}
114 def getComponentsToCompute(self, r, component):
115 """ Returns the list of components necessary to compute component.
117 :param r: Radius of computation
118 :type r: float
119 :param component: Component to compute
120 :type component: str
121 :return: First one lists the components to compute, \
122 second one the components that are computing (cache is created but is None)
123 :rtype: list, list
124 """
126 # Initialise the cached data
127 if r not in self.cachedRThetaPhiData:
128 self.cachedRThetaPhiData[r] = {}
130 components = [component]
131 # Check if additional components are needed (r and th and ph for norm)
132 if component == "norm":
133 for additionalComponent in self.components:
134 components.append(additionalComponent)
136 toCompute = []
137 computing = []
139 # Check if the computation states of the asked components
140 for c in components:
141 # If no cached data exists, add c to the components to compute
142 if c not in self.cachedRThetaPhiData[r]:
143 toCompute.append(c)
144 # If cached data exists but is None it means that c is being computed
145 elif self.cachedRThetaPhiData[r][c] is None:
146 computing.append(c)
148 return toCompute, computing
150 def isAlreadyComputed(self, r, component):
151 """ Tells if a component is to be computed or being computed.
153 :param r: Radius of computation
154 :type r: float
155 :param component: Component to check
156 :type component: str
157 :return: First bool tells if the component is to be computed, \
158 second one tells if the component is being computed
159 :rtype: bool,bool
160 """
161 toCompute, computing = self.getComponentsToCompute(r, component)
162 return len(toCompute) > 0, len(computing) > 0
164 def getRThetaPhiGridData(self, r, itime, removemean, component):
165 """
166 Computes the data on the angular grid for a given grid and caches it.
168 :param r: radius of computation
169 :type r: float
170 :param itime: index of the computation date
171 :type itime: int
172 :param removemean: indicating if the mean should be removed when computing
173 :type removemean: bool
174 :param component: component to compute
175 :type component: str
176 :return: computed data
177 :rtype: np.array
178 """
179 if not self.has_grid:
180 raise ValueError("Can not get grid data, data has no defined real space (th,ph) grid")
181 toCompute, computing = self.getComponentsToCompute(r, component)
183 for component in toCompute:
184 self.cachedRThetaPhiData[r][component] = None
186 if len(toCompute) > 0:
187 result = self.computeRThetaPhiData(r, self.th, self.ph, toCompute, usegridPnm=True, computeallrealisation=False, irealisation=-1)
188 for component in toCompute:
189 print("Preparing to cache data for ", component) # DEBUG
190 self.cachedRThetaPhiData[r][component] = result[component]
191 if(component == "norm"): print(result[component])
193 self.calculateTemporalMean(r, toCompute)
195 if removemean:
196 if component == "norm":
197 return self.computeNorm(r, itime)
198 else:
199 return self.cachedRThetaPhiData[r][component][itime] - self.cachedRThetaPhiData_mean[r][component]
200 else:
201 return self.cachedRThetaPhiData[r][component][itime]
203 def computeSpectra(self, spectraType, spectraDate, itime):
204 self.computingState = 0
206 if spectraType == "spectraofmean":
207 # Do ensemble AND temporal mean
208 if self.has_realisations:
209 mean_gauss_coefs = self.data.mean(axis=(0, 2))
210 # Do temporal mean if no realisations
211 else:
212 mean_gauss_coefs = self.data.mean(axis=0)
214 self.computingState = 50
215 # Compute the spectra of the mean
216 degrees, spectrum_to_return = self.computeLowes(mean_gauss_coefs, squared=False)
218 elif spectraType == "meanofspectra":
219 squared_data = np.square(self.data)
221 if self.has_realisations:
222 degrees, full_lowes_spectrum = self.computeLowes(np.moveaxis(squared_data, 1, 2), squared=True)
224 # Do temporal and ensemble mean of the computed spectra
225 spectrum_to_return = full_lowes_spectrum.mean(axis=(0, 1))
226 else:
227 degrees, full_lowes_spectrum = self.computeLowes(squared_data, squared=True)
229 # Do temporal mean of the computed spectra
230 spectrum_to_return = full_lowes_spectrum.mean(axis=0)
232 elif spectraType == "deviation":
233 if self.has_realisations:
234 mean_gauss_coefs = self.data.mean(axis=2)
235 else:
236 raise ValueError('Cannot compute deviation if the measure has no realisation.')
238 sq_deviation_gauss = np.square(self.data - mean_gauss_coefs[..., np.newaxis])
239 degrees, full_lowes_spectrum = self.computeLowes(np.moveaxis(sq_deviation_gauss, 1, 2), squared=True)
241 # Do temporal and ensemble mean of the computed spectra
242 spectrum_to_return = full_lowes_spectrum.mean(axis=(0, 1))
244 elif spectraType == "dated":
246 if self.has_realisations:
247 gauss_coefs = self.data.mean(axis=2)
248 else:
249 gauss_coefs = self.data
251 self.computingState = 50
252 degrees, spectrum_to_return = self.computeLowes(gauss_coefs[itime], squared=False)
254 else:
255 raise ValueError('Spectra type not understood')
257 self.computingState = 100
259 return degrees, spectrum_to_return
261 def computeLowes(self, gauss_coefs, squared):
262 """
263 Computes the Lowes spectrum from the gauss coefficients given. Handles broadcasting as long as the last dimension is the number of gauss coefficients
265 :param gauss_coefs: NumPy array containing the (squared or not) gauss coeffs for every degree. The last dimension must be the number of gauss coefficients.
266 :type gauss_coefs: np.array (dim: ..., nb_coeffs)
267 :param squared: indicates if the supplied gauss_coefs are already squared or not
268 :type squared: bool
269 :return: NumPy array containing the degrees, NumPy array containing the Lowes spectra
270 :rtype: np.array(dim: lmax), np.array (dim: ..., lmax)
271 """
272 if not squared:
273 squared_gauss_coefs = np.square(gauss_coefs)
274 else:
275 squared_gauss_coefs = gauss_coefs
277 # For broadcast, lowes spectrum must have the same shape as the coeffs except for the last dimension that is lmax.
278 shape_lowes_sp = list(squared_gauss_coefs.shape)
279 shape_lowes_sp[-1] = self.lmax
280 lowes_spectrum = np.zeros(shape=shape_lowes_sp)
282 for l in range(0, self.lmax):
283 n_min = l*(l+2)
284 n_max = (l+1)*(l+3)
285 # Indexes of _LowesPrefactor is the degree so l+1 here (l goes from 0 to lmax - 1)
286 lowes_spectrum[..., l] = self._LowesPrefactor(l+1)*np.sum(squared_gauss_coefs[..., n_min:n_max], axis=-1)
288 return np.arange(1, self.lmax+1, 1), lowes_spectrum
290 def calculateTemporalMean(self, r, components):
291 """
292 Computes the temporal mean and stores it in the cache.
294 :param r: radius where to compute the norm
295 :type r: float
296 :param components: list of components for which the norm should be comuputed
297 :type components: list
298 """
299 if r not in self.cachedRThetaPhiData_mean:
300 self.cachedRThetaPhiData_mean[r] = {}
302 for component in components:
303 self.cachedRThetaPhiData_mean[r][component] = np.mean(self.cachedRThetaPhiData[r][component], axis=0)
305 def expand_dims(self, v):
306 """
307 Returns a copy of v with expanded dimensions to fit realisations
309 :param v: Numpy array to expand
310 :type v: np.ndarray (Ndim)
311 :return: np.ndarray (1 x Ndim x 1)
312 """
313 return np.expand_dims(np.expand_dims(v.copy(), 1), 0)