1 | # -*- coding: utf-8 -*- |
---|
2 | ## =========================================================================== |
---|
3 | ## |
---|
4 | ## This software is governed by the CeCILL license under French law and |
---|
5 | ## abiding by the rules of distribution of free software. You can use, |
---|
6 | ## modify and/ or redistribute the software under the terms of the CeCILL |
---|
7 | ## license as circulated by CEA, CNRS and INRIA at the following URL |
---|
8 | ## "http://www.cecill.info". |
---|
9 | ## |
---|
10 | ## Warning, to install, configure, run, use any of Olivier Marti's |
---|
11 | ## software or to read the associated documentation you'll need at least |
---|
12 | ## one (1) brain in a reasonably working order. Lack of this implement |
---|
13 | ## will void any warranties (either express or implied). |
---|
14 | ## O. Marti assumes no responsability for errors, omissions, |
---|
15 | ## data loss, or any other consequences caused directly or indirectly by |
---|
16 | ## the usage of his software by incorrectly or partially configured |
---|
17 | ## personal. |
---|
18 | ## |
---|
19 | ## =========================================================================== |
---|
20 | ''' |
---|
21 | Utilities for LMDZ grid |
---|
22 | |
---|
23 | olivier.marti@lsce.ipsl.fr |
---|
24 | ''' |
---|
25 | |
---|
26 | ## SVN information |
---|
27 | __Author__ = "$Author$" |
---|
28 | __Date__ = "$Date$" |
---|
29 | __Revision__ = "$Revision$" |
---|
30 | __Id__ = "$Id: $" |
---|
31 | __HeadURL = "$HeadURL$" |
---|
32 | |
---|
33 | import numpy as np |
---|
34 | |
---|
35 | try : import xarray as xr |
---|
36 | except ImportError : pass |
---|
37 | |
---|
38 | try : import numba |
---|
39 | except ImportError : pass |
---|
40 | |
---|
41 | rpi = np.pi ; rad = np.deg2rad (1.0) ; dar = np.rad2deg (1.0) |
---|
42 | |
---|
43 | |
---|
44 | def __mmath__ (tab) : |
---|
45 | ''' |
---|
46 | Determines the type of tab : i.e. numpy or xarray |
---|
47 | ''' |
---|
48 | math = None |
---|
49 | try : |
---|
50 | if type (tab) == xr.core.dataarray.DataArray : math = xr |
---|
51 | except : |
---|
52 | pass |
---|
53 | |
---|
54 | try : |
---|
55 | if type (tab) == np.ndarray : math = np |
---|
56 | except : |
---|
57 | pass |
---|
58 | |
---|
59 | return math |
---|
60 | # |
---|
61 | def extend (tab, Lon=False, jplus=25, jpi=None, lonplus=360.0) : |
---|
62 | ''' |
---|
63 | Returns extended field eastward to have better plots, and box average crossing the boundary |
---|
64 | Works only for xarray and numpy data (?) |
---|
65 | |
---|
66 | tab : field to extend. |
---|
67 | Lon : (optional, default=False) : if True, add 360 in the extended parts of the field |
---|
68 | jpi : normal longitude dimension of the field. exrtend does nothing it the actual |
---|
69 | size of the field != jpi (avoid to extend several times) |
---|
70 | jplus (optional, default=25) : number of points added on the east side of the field |
---|
71 | |
---|
72 | ''' |
---|
73 | math = __mmath__ (tab) |
---|
74 | if tab.shape[-1] == 1 : extend = tab |
---|
75 | |
---|
76 | else : |
---|
77 | if jpi == None : jpi = tab.shape[-1] |
---|
78 | |
---|
79 | if Lon : xplus = lonplus |
---|
80 | else : xplus = 0.0 |
---|
81 | |
---|
82 | if tab.shape[-1] > jpi : |
---|
83 | extend = tab |
---|
84 | else : |
---|
85 | istart = 0 ; le=jpi+1 ; la=0 |
---|
86 | if math == xr : |
---|
87 | lon = tab.dims[-1] |
---|
88 | extend = xr.concat ((tab.isel(lon=slice(istart ,istart+le )), |
---|
89 | tab.isel(lon=slice(istart+la,istart+la+jplus))+xplus ), dim=lon) |
---|
90 | try : |
---|
91 | extend_lon = xr.concat ((tab.coords[lon].isel(lon=slice(istart ,istart+le )), |
---|
92 | tab.coords[lon].isel(lon=slice(istart+la,istart+la+jplus))+lonplus), dim=lon) |
---|
93 | extend = extend.assign_coords ( {tab.dims[-1]:extend_lon} ) |
---|
94 | except : |
---|
95 | pass |
---|
96 | if math == np : |
---|
97 | extend = np.concatenate ((tab [..., istart:istart+le], tab [..., istart+la:jplus]+xplus ), axis=-1) |
---|
98 | |
---|
99 | return extend |
---|
100 | |
---|
101 | def interp1d (x, xp, yp, zdim='presnivs', units=None, verbose=False, method='linear') : |
---|
102 | ''' |
---|
103 | One-dimensionnal interpolation of a multi-dimensionnal field |
---|
104 | |
---|
105 | Intended to interpolate on standard pressure level |
---|
106 | |
---|
107 | All inputs shoud be xarray data arrays |
---|
108 | |
---|
109 | Input : |
---|
110 | x : levels at wich we want to interpolate (i.e. standard pressure levels) |
---|
111 | xp : position of the input points (i.e. pressure) |
---|
112 | yp : fields values at these points (temperature, humidity, etc ..) |
---|
113 | zdim : name of the dimension that we want to interpolate |
---|
114 | method : available methods are |
---|
115 | linear |
---|
116 | log[arithmic]. Available for positive values only |
---|
117 | nearest : nearest neighbor |
---|
118 | |
---|
119 | \!/ xp should be decreasing values along zdim axis \!/ |
---|
120 | ''' |
---|
121 | # Get the number of dimension with dim==zdim |
---|
122 | axis = list(xp.dims).index(zdim) |
---|
123 | |
---|
124 | # Get the number of levels in each arrays |
---|
125 | nk_in = xp.shape[axis] |
---|
126 | nk_ou = len (x) |
---|
127 | |
---|
128 | # Define the result array |
---|
129 | in_shape = np.array (xp.shape) |
---|
130 | if verbose : print ( f'{in_shape=}' ) |
---|
131 | ou_shape = np.array (in_shape) |
---|
132 | if verbose : print ( f'{ou_shape=}' ) |
---|
133 | ou_shape[axis] = nk_ou |
---|
134 | |
---|
135 | in_dims = list (yp.dims) |
---|
136 | if verbose : print ( f'{in_dims=}' ) |
---|
137 | ou_dims = in_dims |
---|
138 | |
---|
139 | pdim = x.dims[0] |
---|
140 | ou_dims[axis] = pdim |
---|
141 | |
---|
142 | new_coords = [] |
---|
143 | for i, dim in enumerate (yp.dims) : |
---|
144 | if dim == zdim : |
---|
145 | ou_dims[i] = x.dims[0] |
---|
146 | if units != None : yp[dim].attrs['units'] = units |
---|
147 | new_coords.append (x .values) |
---|
148 | else : |
---|
149 | new_coords.append (yp.coords[dim].values) |
---|
150 | |
---|
151 | if verbose : |
---|
152 | print ( f'{ou_dims =}' ) |
---|
153 | print ( f'{new_coords=}' ) |
---|
154 | |
---|
155 | ou_tab = xr.DataArray (np.empty (ou_shape), dims=ou_dims, coords=new_coords) |
---|
156 | |
---|
157 | if 'log' in method : |
---|
158 | yp_min = yp.min() ; yp_max = yp.max() |
---|
159 | indic = yp_min * yp_max |
---|
160 | if indic <= 0. : |
---|
161 | print ('Input data have a change of sign') |
---|
162 | print ('Error : logarithmic method is available only for') |
---|
163 | print ('positive or negative input values. ') |
---|
164 | raise Exception |
---|
165 | |
---|
166 | # Optimized (pre-compiled) interpolation loop |
---|
167 | #@numba.jit (nopython=True) |
---|
168 | def __interp (nk_ou, x, xp, yp) : |
---|
169 | # Interpolate |
---|
170 | # Find index of the just above level |
---|
171 | idk1 = np.minimum ( (x-xp), 0.).argmax (dim=zdim) |
---|
172 | idk2 = idk1 - 1 |
---|
173 | idk2 = np.maximum (idk2, 0) |
---|
174 | |
---|
175 | x1 = xp[{zdim:idk1}] |
---|
176 | x2 = xp[{zdim:idk2}] |
---|
177 | |
---|
178 | dx1 = x - x1 |
---|
179 | dx2 = x2 - x |
---|
180 | dx = x2 - x1 |
---|
181 | dx1 = dx1/dx ; dx2 = dx2/dx |
---|
182 | |
---|
183 | y1 = yp[{zdim:idk1}] |
---|
184 | y2 = yp[{zdim:idk2}] |
---|
185 | |
---|
186 | #print ( f'{k=} {idk1=} {idk2=} {x1=} {x2=} {dx=1} {dx2=} {y1=} {y2}' ) |
---|
187 | |
---|
188 | if 'linear' in method : |
---|
189 | result = (dx1*y2 + dx2*y1) |
---|
190 | if 'log' in method : |
---|
191 | if yp_min > 0. : |
---|
192 | result = np.power(y2, dx1) * np.power(y1, dx2) |
---|
193 | if yp_max < 0. : |
---|
194 | result = -np.power(-y2, dx1) * np.power(-y1, dx2) |
---|
195 | if 'nearest' in method : |
---|
196 | result = xr.where ( dx2>=dx1, y1, y2) |
---|
197 | |
---|
198 | return result |
---|
199 | |
---|
200 | for k in np.arange (nk_ou) : |
---|
201 | result = __interp (nk_ou, x[{pdim:k}], xp, yp) |
---|
202 | |
---|
203 | # Put result in the final array |
---|
204 | ou_tab [{pdim:k}] = result |
---|
205 | |
---|
206 | return ou_tab.squeeze() |
---|
207 | |
---|
208 | def fixed_lon (lon, center_lon=0.0) : |
---|
209 | ''' |
---|
210 | Returns corrected longitudes for nicer plots |
---|
211 | |
---|
212 | lon : longitudes of the grid. At least 1D. |
---|
213 | center_lon : center longitude. Default=0. |
---|
214 | |
---|
215 | Designed by Phil Pelson. See https://gist.github.com/pelson/79cf31ef324774c97ae7 |
---|
216 | ''' |
---|
217 | mmath = __mmath__ (lon) |
---|
218 | |
---|
219 | fixed_lon = lon.copy () |
---|
220 | |
---|
221 | fixed_lon = mmath.where (fixed_lon > center_lon+180., fixed_lon-360.0, fixed_lon) |
---|
222 | fixed_lon = mmath.where (fixed_lon < center_lon-180., fixed_lon+360.0, fixed_lon) |
---|
223 | |
---|
224 | start = np.argmax (np.abs (np.diff (fixed_lon, axis=-1)) > 180., axis=-1) |
---|
225 | fixed_lon [start+1:] += 360. |
---|
226 | |
---|
227 | return fixed_lon |
---|
228 | |
---|
229 | def nord2sud (p2D) : |
---|
230 | ''' |
---|
231 | Swap north to south a 2D field |
---|
232 | ''' |
---|
233 | pout = p2D [..., -1::-1, : ] |
---|
234 | |
---|
235 | return pout |
---|
236 | |
---|
237 | def point2geo (p1D) : |
---|
238 | ''' |
---|
239 | From 1D (restart type) to 2D |
---|
240 | ''' |
---|
241 | math = __mmath__ (p1D) |
---|
242 | |
---|
243 | # Get the dimensions |
---|
244 | jpn = p1D.shape[-1] |
---|
245 | |
---|
246 | if len (p1D.shape) > 1 : |
---|
247 | jpt = p1D.shape[0] |
---|
248 | else : |
---|
249 | jpt = 0 |
---|
250 | |
---|
251 | if jpn == 9026 : jpi = 96 ; jpj = 96 |
---|
252 | if jpn == 17858 : jpi = 144 ; jpj = 144 |
---|
253 | if jpn == 20306 : jpi = 144 ; jpj = 143 |
---|
254 | |
---|
255 | if jpt > 0 : |
---|
256 | p2D = np.zeros ( (jpt, jpj, jpi) ) |
---|
257 | p2D [:, 1:jpj-1, :] = np.reshape (p1D [:,1:jpn-1], (jpt, jpj-2, jpi) ) |
---|
258 | p2D [:, 0 , : ] = p1D[:, 0 ] |
---|
259 | p2D [:, jpj-1, : ] = p1D[:, jpn-1] |
---|
260 | |
---|
261 | else : |
---|
262 | p2D = np.zeros ( (jpj, jpi) ) |
---|
263 | p2D [1:jpj-1, :] = np.reshape (np.float64 (p1D [1:jpn-1]), (jpj-2, jpi) ) |
---|
264 | p2D [0 , : ] = p1D[ 0 ] |
---|
265 | p2D [jpj-1, : ] = p1D[jpn-1] |
---|
266 | |
---|
267 | if math == xr : |
---|
268 | p2D = xr.DataArray (p2D) |
---|
269 | for attr in p1D.attrs : |
---|
270 | p2D.attrs[attr] = p1D.attrs[attr] |
---|
271 | p2D = p2D.rename ( {p2D.dims[0]:p1D.dims[0], p2D.dims[-1]:'x', p2D.dims[-2]:'y'} ) |
---|
272 | |
---|
273 | return p2D |
---|
274 | |
---|
275 | def geo2point ( p2D, cumulPoles=False, dim1D='points_physiques' ) : |
---|
276 | ''' |
---|
277 | From 2D (lat, lon) to 1D (points_phyiques) |
---|
278 | ''' |
---|
279 | math = __mmath__ (p2D) |
---|
280 | # |
---|
281 | # Get the dimension |
---|
282 | |
---|
283 | (jpj, jpi) = p2D.shape[-2:] |
---|
284 | |
---|
285 | if len (p2D.shape) > 2 : |
---|
286 | jpt = p2D.shape[0] |
---|
287 | else : |
---|
288 | jpt = -1 |
---|
289 | |
---|
290 | jpn = jpi*(jpj-2) + 2 |
---|
291 | |
---|
292 | if jpt == -1 : |
---|
293 | p1D = np.zeros ( (jpn,) ) |
---|
294 | p1D[1:-1] = np.float64(p2D[1:-1, :]).ravel() |
---|
295 | p1D[ 0] = p2D[ 0, 0] |
---|
296 | p1D[-1] = p2D[-1, 0] |
---|
297 | |
---|
298 | else : |
---|
299 | p1D = np.zeros ( (jpt, jpn) ) |
---|
300 | if math == xr : |
---|
301 | p1D[:, 1:-1] = np.reshape ( np.float64 (p2D[:, 1:-1, :].values).ravel(), (jpt, jpn-2) ) |
---|
302 | else : |
---|
303 | p1D[:, 1:-1] = np.reshape ( np.float64 (p2D[:, 1:-1, :] ).ravel(), (jpt, jpn-2) ) |
---|
304 | p1D[:, 0 ] = p2D[:, 0, 0] |
---|
305 | p1D[:, -1 ] = p2D[:, -1, 0] |
---|
306 | |
---|
307 | if math == xr : |
---|
308 | p1D = xr.DataArray (p1D) |
---|
309 | for attr in p2D.attrs : |
---|
310 | p1D.attrs[attr] = p2D.attrs[attr] |
---|
311 | p1D = p1D.rename ( {p1D.dims[0]:p2D.dims[0], p1D.dims[-1]:dim1D} ) |
---|
312 | |
---|
313 | if cumulPoles : |
---|
314 | p1D[..., 0] = np.sum ( p2D[..., 0, :] ) |
---|
315 | p1D[..., -1] = np.sum ( p2D[..., -1, :] ) |
---|
316 | |
---|
317 | return p1D |
---|
318 | |
---|
319 | def geo3point ( p3D, cumulPoles=False, dim1D='points_physiques' ) : |
---|
320 | ''' |
---|
321 | From 3D (lev, lat, lon) to 2D (lev, points_phyiques) |
---|
322 | ''' |
---|
323 | math = __mmath__ (p3D) |
---|
324 | # |
---|
325 | # Get the dimensions |
---|
326 | |
---|
327 | (jpk, jpj, jpi) = p3D.shape[-3:] |
---|
328 | |
---|
329 | if len (p3D.shape) > 3 : |
---|
330 | jpt = p3D.shape[0] |
---|
331 | else : |
---|
332 | jpt = -1 |
---|
333 | |
---|
334 | jpn = jpi*(jpj-2) + 2 |
---|
335 | |
---|
336 | if jpt == -1 : |
---|
337 | p2D = np.zeros ( (jpk, jpn,) ) |
---|
338 | for jk in np.arange (jpk) : |
---|
339 | p2D [jk, :] = geo2point ( p3D [jk,:,:], cumulPoles, dim1D ) |
---|
340 | else : |
---|
341 | p2D = np.zeros ( (jpt, jpk, jpn) ) |
---|
342 | for jk in np.arange (jpk) : |
---|
343 | p2D [:, jk, :] = geo2point ( p3D [:, jk,:,:], cumulPoles, dim1D ) |
---|
344 | |
---|
345 | if math == xr : |
---|
346 | p2D = xr.DataArray (p2D) |
---|
347 | for attr in p2D.attrs : |
---|
348 | p2D.attrs[attr] = p3D.attrs[attr] |
---|
349 | p2D = p2D.rename ( {p2D.dims[-1]:dim1D, p2D.dims[-2]:p3D.dims[-3]} ) |
---|
350 | |
---|
351 | return p2D |
---|
352 | |
---|
353 | def geo2en (pxx, pyy, pzz, glam, gphi) : |
---|
354 | ''' |
---|
355 | Change vector from geocentric to east/north |
---|
356 | |
---|
357 | Inputs : |
---|
358 | pxx, pyy, pzz : components on the geocentric system |
---|
359 | glam, gphi : longitude and latitude of the points |
---|
360 | ''' |
---|
361 | |
---|
362 | gsinlon = np.sin (rad * glam) |
---|
363 | gcoslon = np.cos (rad * glam) |
---|
364 | gsinlat = np.sin (rad * gphi) |
---|
365 | gcoslat = np.cos (rad * gphi) |
---|
366 | |
---|
367 | pte = - pxx * gsinlon + pyy * gcoslon |
---|
368 | ptn = - pxx * gcoslon * gsinlat - pyy * gsinlon * gsinlat + pzz * gcoslat |
---|
369 | |
---|
370 | return pte, ptn |
---|
371 | |
---|
372 | def en2geo (pte, ptn, glam, gphi) : |
---|
373 | ''' |
---|
374 | Change vector from east/north to geocentric |
---|
375 | |
---|
376 | Inputs : |
---|
377 | pte, ptn : eastward/northward components |
---|
378 | glam, gphi : longitude and latitude of the points |
---|
379 | ''' |
---|
380 | |
---|
381 | gsinlon = np.sin (rad * glam) |
---|
382 | gcoslon = np.cos (rad * glam) |
---|
383 | gsinlat = np.sin (rad * gphi) |
---|
384 | gcoslat = np.cos (rad * gphi) |
---|
385 | |
---|
386 | pxx = - pte * gsinlon - ptn * gcoslon * gsinlat |
---|
387 | pyy = pte * gcoslon - ptn * gsinlon * gsinlat |
---|
388 | pzz = ptn * gcoslat |
---|
389 | |
---|
390 | return pxx, pyy, pzz |
---|
391 | |
---|
392 | ## =========================================================================== |
---|
393 | ## |
---|
394 | ## That's all folk's !!! |
---|
395 | ## |
---|
396 | ## =========================================================================== |
---|