Coverage for webgeodyn/data/GHData.py: 0%
123 statements
« prev ^ index » next coverage.py v7.2.7, created at 2023-12-18 09:33 +0000
1import numpy as np
2from .measuredata import MeasureData
3from .normPnm import semiNormalisedSchmidt
4from webgeodyn.constants import rE
7class GHData(MeasureData):
8 """ Class handling MF or SV, DIFF, ER data (g and h harmonics coefficients)."""
9 def __init__(self, data, lmax, units, measureType, PnmNorm=semiNormalisedSchmidt, gridth=None, gridph=None):
10 """ Constructor of GHData class
11 :param data: measure data
12 :type data: np.array
13 :param lmax: maximum degree
14 :type lmax: int
15 :param units: units of the measure
16 :type units: str
17 :param measureType: type of the measure
18 :type measureType: str ('SV', 'MF' or 'LOD')
19 :param PnmNorm: normalisation to use (default: semiNormalisedSchmidt)
20 :type PnmNorm: function
21 :param gridth: grid on which the theta are evaluated
22 :type gridth: np.array
23 :param gridph: grid on which the theta are evaluated
24 :type gridph: np.array
25 """
26 super().__init__(data.shape, lmax, PnmNorm, units, gridth, gridph)
27 self.measureType = measureType
28 self.initMeasure()
29 self.setData(data)
31 def initMeasure(self):
32 """ Initiates the measure parameters : components, coefs and grids if needed. Called in __init__. """
33 self.nk = self.lmax * (self.lmax + 2)
34 self.components = ["r", "th", "ph"]
35 # Only one coef for LOD
36 if self.measureType == "LOD":
37 self.coefs_list = ["LOD"]
38 elif self.measureType == "EXT_DIPOLE":
39 self.coefs_list = ["q", "s"]
40 else:
41 self.coefs_list = ["g", "h"]
42 self.initMatrices()
43 if self.th is not None and self.ph is not None:
44 self.computeGridPnm()
46 def setData(self, data):
47 """ Sets the input data and extracts the time array and temporal mean of the coefs. Called in __init__.
48 :param data: measure data
49 :type data: np.array
50 """
51 # Clear cached data to avoid keeping old data
52 self.clearCache()
53 self.data = data[:, :self.nk]
54 self.ntimes = data.shape[0]
55 self.coefs_mean = np.mean(self.data, axis=0)
56 self.has_realisations = False
57 if self.data.ndim >= 3:
58 self.has_realisations = True
60 def lm2k(self, l, m, coef):
61 """ Computes the index k of a coef given its name, its degree l and its order m. Returns None if no coef exists.
62 :param l: spherical harmonic degree
63 :type l: int
64 :param m: spherical harmonic order
65 :type m: int
66 :param coef: type of coef
67 :type coef: str ('g', 'h' or 'LOD')
68 :return: index of the coef in data
69 :rtype: int or None
70 """
71 if m > l:
72 return None
73 if m == 0:
74 if coef == "h" or coef == "s":
75 return None
76 return int(l ** 2 - 1)
77 if coef == "g" or coef == "q":
78 return int(l ** 2 + 2 * m - 2)
79 if coef == "h" or coef == "s":
80 return int(l ** 2 + 2 * m - 1)
81 if coef == "LOD":
82 return 0
84 def k2lm(self, k):
85 """ Returns the degree, the order and the name of a coef of given k.
86 :param k: index of the coef
87 :type k: int
88 :return: degree, order and name of the coef of index k
89 :rtype: int, int, str
90 """
91 return self.l[k], self.m[k], self.coefname[k]
93 def initMatrices(self):
94 """ Creates the matrices that allow to identify and index the spherical
95 harmonic coefficients g and h. Called in __init__.
96 """
97 k = np.arange(self.lmax * (self.lmax + 2))
98 self.l = np.floor(np.sqrt(k + 1)).astype(int)
99 kl2 = (k + 1 - self.l ** 2)
100 self.is_h = np.logical_and(np.mod(kl2, 2) == 0, kl2 != 0)
101 self.is_g = np.logical_not(self.is_h)
102 self.coefname = np.where(self.is_g, "g", "h")
103 self.m = np.ceil(kl2 / 2).astype(int)
105 def set_expand_dims(self, useRealisations):
106 """
107 Expands the dims of degrees and orders to be able to broadcast realisations. Also creates the l+1 array.
109 :param useRealisations: expands only if this is True
110 :type useRealisations: bool
111 """
112 if useRealisations:
113 # for broadcast dimension (coef is the second dimension on 3)
114 self.is_h_ = self.expand_dims(self.is_h)
115 self.is_g_ = self.expand_dims(self.is_g)
116 self.m_ = self.expand_dims(self.m)
117 self.l_ = self.expand_dims(self.l)
118 else:
119 self.is_h_ = self.is_h
120 self.is_g_ = self.is_g
121 self.m_ = self.m
122 self.l_ = self.l
123 self.lplus1 = self.l_ + 1
125 def computeRThetaPhiData(self, r, thlist, phlist, components=None, usegridPnm=False, computeallrealisation=False,
126 irealisation=-1):
127 """ Computes the components of the quantity described by the spherical harmonic potential
128 given radius, a colatitude list and a longitude list.
130 :param r: radius at which the computation is done
131 :type r: float
132 :param thlist: list containing the colatitudes at which the computation is done
133 :type thlist: list
134 :param phlist: list containing the longitudes at which the computation is done
135 :type phlist: list
136 :param components: list of components to compute (default to ["r", "th", "ph", "norm"])
137 :type components: list
138 :param usegridPnm: boolean indicating if the Pnm computed previously should be used (default=False)
139 :type usegridPnm: bool
140 :param computeallrealisation: boolean indicating if the computation should be on all realisations or on the mean (default=False)
141 :type computeallrealisation: bool
142 :param irealisation: index of the realisation to use (default=-1 and takes the mean instead)
143 :type irealisation: int
144 :return: result of the computation
145 :rtype: dict of np.array of dim (date x n_ph x n_th)
146 """
147 # Default value of components
148 if components is None:
149 components = ["r", "th", "ph", "norm"]
151 if ("norm" in components) and not usegridPnm: # if norm, require th and ph component
152 for component in self.components:
153 if component not in components:
154 components.append(component)
156 if self.has_realisations and not computeallrealisation:
157 if irealisation < 0:
158 realisation_data = np.mean(self.data, axis=2)
159 else:
160 realisation_data = self.data[:, :, irealisation]
161 else:
162 realisation_data = self.data
164 if self.has_realisations and computeallrealisation:
165 result = {component: np.zeros((self.ntimes, len(thlist), len(phlist), realisation_data.shape[2]))
166 for component in components}
167 # Use expanded arrays to fit dimension of realisations
168 self.set_expand_dims(True)
169 else:
170 result = {component: np.zeros((self.ntimes, len(thlist), len(phlist))) for component in components}
171 self.set_expand_dims(False)
173 # Compute a handy prefactor before looping
174 rEr_lplus2 = (rE / r) ** (self.l_ + 2)
176 # Loop on the colatitude list
177 for ith, th in enumerate(thlist):
178 if usegridPnm:
179 if not self.has_grid:
180 raise ValueError("The data has no real space (th,ph) grid")
181 self.computingState = (ith + 1) / self.nth * 100
182 sinth = np.sin(th)
183 if usegridPnm:
184 Pnm = self.Pnm[ith]
185 dPnm = self.dPnm[ith]
186 else:
187 Pnm, dPnm = self.computePnm(th)
189 if self.has_realisations and computeallrealisation:
190 Pnm = self.expand_dims(Pnm)
191 dPnm = self.expand_dims(dPnm)
193 # Loop on the longitude list
194 for iph, ph in enumerate(phlist):
195 # Calculate the radial component
196 if "r" in components:
197 fac_gnm_r = self.lplus1 * rEr_lplus2 * np.cos(self.m_ * ph) * self.is_g_
198 fac_hnm_r = self.lplus1 * rEr_lplus2 * np.sin(self.m_ * ph) * self.is_h_
199 result["r"][:, ith, iph] = np.sum(
200 ((fac_gnm_r + fac_hnm_r) * realisation_data) * Pnm,
201 axis=1)
203 # Calculate the component along theta
204 if "th" in components:
205 fac_gnm_th = - rEr_lplus2 * np.cos(self.m_ * ph)
206 fac_hnm_th = - rEr_lplus2 * np.sin(self.m_ * ph)
207 result["th"][:, ith, iph] = np.sum(
208 (fac_gnm_th * (realisation_data * self.is_g_)
209 + fac_hnm_th * (realisation_data * self.is_h_))
210 * dPnm,
211 axis=1)
213 # Calculate the component along phi
214 if "ph" in components:
215 fac_gnm_ph = rEr_lplus2 * self.m_ * np.sin(self.m_ * ph) / sinth
216 fac_hnm_ph = - rEr_lplus2 * self.m_ * np.cos(self.m_ * ph) / sinth
217 result["ph"][:, ith, iph] = np.sum(
218 (fac_gnm_ph * (realisation_data * self.is_g_)
219 + fac_hnm_ph * (realisation_data * self.is_h_))
220 * Pnm,
221 axis=1)
223 # Calculate the norm of the obtained vector
224 if "norm" in components:
225 if "r" in result:
226 resultr = result["r"][:, ith, iph]
227 else:
228 resultr = self.cachedRThetaPhiData[r]["r"][:, ith, iph]
230 if "th" in result:
231 resultth = result["th"][:, ith, iph]
232 else:
233 resultth = self.cachedRThetaPhiData[r]["th"][:, ith, iph]
235 if "ph" in result:
236 resultph = result["ph"][:, ith, iph]
237 else:
238 resultph = self.cachedRThetaPhiData[r]["ph"][:, ith, iph]
240 result["norm"][:, ith, iph] = np.sqrt(resultr ** 2 + resultth ** 2 + resultph ** 2)
242 # End of angular loops
244 return result
246 def _LowesPrefactor(self, l):
247 return l + 1
249 def computeNorm(self, r, itime):
250 """
251 Computes norm at a certain time and place using cached data.
253 :param r: radius at which to evaluate the norm
254 :type r: float
255 :param itime: index of time
256 :type itime: int
257 :return: Numpy array
258 :rtype: np.array
259 """
260 return np.sqrt((self.cachedRThetaPhiData[r]["th"][itime] - self.cachedRThetaPhiData_mean[r]["th"]) ** 2
261 + (self.cachedRThetaPhiData[r]["ph"][itime] - self.cachedRThetaPhiData_mean[r]["ph"]) ** 2
262 + (self.cachedRThetaPhiData[r]["r"][itime] - self.cachedRThetaPhiData_mean[r]["r"]) ** 2)