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 | - Lots of tests for xarray object |
---|
24 | - Not much tested for numpy objects |
---|
25 | |
---|
26 | Author: olivier.marti@lsce.ipsl.fr |
---|
27 | |
---|
28 | ## SVN information |
---|
29 | __Author__ = "$Author$" |
---|
30 | __Date__ = "$Date$" |
---|
31 | __Revision__ = "$Revision$" |
---|
32 | __Id__ = "$Id: $" |
---|
33 | __HeadURL = "$HeadURL$" |
---|
34 | ''' |
---|
35 | |
---|
36 | import numpy as np |
---|
37 | import xarray as xr |
---|
38 | |
---|
39 | RPI = np.pi |
---|
40 | RAD = np.deg2rad (1.0) |
---|
41 | DAR = np.rad2deg (1.0) |
---|
42 | REPSI = np.finfo (1.0).eps |
---|
43 | |
---|
44 | RAAMO = 12 # Number of months in one year |
---|
45 | RJJHH = 24 # Number of hours in one day |
---|
46 | RHHMM = 60 # Number of minutes in one hour |
---|
47 | RMMSS = 60 # Number of seconds in one minute |
---|
48 | RA = 6371229.0 # Earth radius [m] |
---|
49 | GRAV = 9.80665 # Gravity [m/s2] |
---|
50 | RT0 = 273.15 # Freezing point of fresh water [Kelvin] |
---|
51 | RAU0 = 1026.0 # Volumic mass of sea water [kg/m3] |
---|
52 | RLEVAP = 2.5e+6 # Latent heat of evaporation (water) [J/K] |
---|
53 | VKARMN = 0.4 # Von Karman constant |
---|
54 | STEFAN = 5.67e-8 # Stefan-Boltzmann constant [W/m2/K4] |
---|
55 | #RHOS = 330. # Volumic mass of snow [kg/m3] |
---|
56 | |
---|
57 | RDAY = RJJHH * RHHMM * RMMSS # Day length [s] |
---|
58 | RSIYEA = 365.25 * RDAY * 2. * RPI / 6.283076 # Sideral year length [s] |
---|
59 | RSIDAY = RDAY / (1. + RDAY / RSIYEA) # Sideral day length [s] |
---|
60 | ROMEGA = 2. * RPI / RSIDAY # Earth rotation parameter [s-1] |
---|
61 | |
---|
62 | def __mmath__ (ptab, default=None) : |
---|
63 | '''Determines the type of tab : xarray, numpy or numpy.ma object ? |
---|
64 | |
---|
65 | Returns type |
---|
66 | ''' |
---|
67 | mmath = default |
---|
68 | if isinstance (ptab, xr.core.dataarray.DataArray) : |
---|
69 | mmath = xr |
---|
70 | if isinstance (ptab, np.ndarray) : |
---|
71 | mmath = np |
---|
72 | if isinstance (ptab, np.ma.MaskType) : |
---|
73 | mmath = np.ma |
---|
74 | |
---|
75 | return mmath |
---|
76 | |
---|
77 | # |
---|
78 | def extend (tab, Lon=False, jplus=25, jpi=None, lonplus=360.0) : |
---|
79 | '''Returns extended field eastward to have better plots, and box average crossing the boundary |
---|
80 | |
---|
81 | Works only for xarray and numpy data (?) |
---|
82 | |
---|
83 | tab : field to extend. |
---|
84 | Lon : (optional, default=False) : if True, add 360 in the extended parts of the field |
---|
85 | jpi : normal longitude dimension of the field. exrtend does nothing it the actual |
---|
86 | size of the field != jpi (avoid to extend several times) |
---|
87 | jplus (optional, default=25) : number of points added on the east side of the field |
---|
88 | ''' |
---|
89 | |
---|
90 | math = __mmath__ (tab) |
---|
91 | if tab.shape[-1] == 1 : |
---|
92 | ztab = tab |
---|
93 | |
---|
94 | else : |
---|
95 | if jpi is None : |
---|
96 | jpi = tab.shape[-1] |
---|
97 | |
---|
98 | if Lon : |
---|
99 | xplus = lonplus |
---|
100 | else : |
---|
101 | xplus = 0.0 |
---|
102 | |
---|
103 | if tab.shape[-1] > jpi : |
---|
104 | ztab = tab |
---|
105 | else : |
---|
106 | istart = 0 |
---|
107 | le = jpi+1 |
---|
108 | la = 0 |
---|
109 | if math == xr : |
---|
110 | lon = tab.dims[-1] |
---|
111 | ztab = xr.concat (( |
---|
112 | tab.isel (lon=slice(istart ,istart+le )), |
---|
113 | tab.isel (lon=slice(istart+la,istart+la+jplus))+xplus ), dim=lon) |
---|
114 | if 'lon' in tab.dims : |
---|
115 | extend_lon = xr.concat (( |
---|
116 | tab.coords[lon].isel(lon=slice(istart ,istart+le )), |
---|
117 | tab.coords[lon].isel(lon=slice(istart+la,istart+la+jplus))+lonplus), dim=lon) |
---|
118 | ztab = ztab.assign_coords ( {tab.dims[-1]:extend_lon} ) |
---|
119 | #except : |
---|
120 | # pass |
---|
121 | if math == np : |
---|
122 | ztab = np.concatenate ((tab [..., istart:istart+le], |
---|
123 | tab [..., istart+la:jplus]+xplus ), axis=-1) |
---|
124 | |
---|
125 | return ztab |
---|
126 | |
---|
127 | def interp1d (x, xp, yp, zdim='presnivs', units=None, verbose=False, method='linear') : |
---|
128 | '''One-dimensionnal interpolation of a multi-dimensionnal field |
---|
129 | |
---|
130 | Intended to interpolate on standard pressure level |
---|
131 | |
---|
132 | All inputs shoud be xarray data arrays |
---|
133 | |
---|
134 | Input : |
---|
135 | x : levels at wich we want to interpolate (i.e. standard pressure levels) |
---|
136 | xp : position of the input points (i.e. pressure) |
---|
137 | yp : fields values at these points (temperature, humidity, etc ..) |
---|
138 | zdim : name of the dimension that we want to interpolate |
---|
139 | method : available methods are |
---|
140 | linear |
---|
141 | log[arithmic]. Available for positive values only |
---|
142 | nearest : nearest neighbor |
---|
143 | |
---|
144 | Warning : xp should be decreasing values along zdim axis |
---|
145 | ''' |
---|
146 | # Get the number of dimension with dim==zdim |
---|
147 | axis = list(xp.dims).index(zdim) |
---|
148 | |
---|
149 | # Get the number of levels in each arrays |
---|
150 | nk_ou = len (x) |
---|
151 | |
---|
152 | # Define the result array |
---|
153 | in_shape = np.array (xp.shape) |
---|
154 | if verbose : |
---|
155 | print ( f'{in_shape=}' ) |
---|
156 | ou_shape = np.array (in_shape) |
---|
157 | if verbose : |
---|
158 | print ( f'{ou_shape=}' ) |
---|
159 | ou_shape[axis] = nk_ou |
---|
160 | |
---|
161 | in_dims = list (yp.dims) |
---|
162 | if verbose : |
---|
163 | print ( f'{in_dims=}' ) |
---|
164 | ou_dims = in_dims |
---|
165 | |
---|
166 | pdim = x.dims[0] |
---|
167 | ou_dims[axis] = pdim |
---|
168 | |
---|
169 | new_coords = [] |
---|
170 | for i, dim in enumerate (yp.dims) : |
---|
171 | if dim == zdim : |
---|
172 | ou_dims[i] = x.dims[0] |
---|
173 | if units is not None : |
---|
174 | yp[dim].attrs['units'] = units |
---|
175 | new_coords.append (x .values) |
---|
176 | else : |
---|
177 | new_coords.append (yp.coords[dim].values) |
---|
178 | |
---|
179 | if verbose : |
---|
180 | print ( f'{ou_dims =}' ) |
---|
181 | print ( f'{new_coords=}' ) |
---|
182 | |
---|
183 | ou_tab = xr.DataArray (np.empty (ou_shape), dims=ou_dims, coords=new_coords) |
---|
184 | |
---|
185 | if 'log' in method : |
---|
186 | yp_min = yp.min() |
---|
187 | yp_max = yp.max() |
---|
188 | indic = yp_min * yp_max |
---|
189 | if indic <= 0. : |
---|
190 | print ( 'Input data have a change of sign') |
---|
191 | print ( 'Error: logarithmic method is available only for') |
---|
192 | print ( 'positive or negative input values. ') |
---|
193 | raise ValueError |
---|
194 | |
---|
195 | # Optimized (pre-compiled) interpolation loop |
---|
196 | #@numba.jit (nopython=True) |
---|
197 | def __interp (x, xp, yp) : |
---|
198 | # Interpolate |
---|
199 | # Find index of the just above level |
---|
200 | idk1 = np.minimum ( (x-xp), 0.).argmax (dim=zdim) |
---|
201 | idk2 = idk1 - 1 |
---|
202 | idk2 = np.maximum (idk2, 0) |
---|
203 | |
---|
204 | x1 = xp[{zdim:idk1}] |
---|
205 | x2 = xp[{zdim:idk2}] |
---|
206 | |
---|
207 | dx1 = x - x1 |
---|
208 | dx2 = x2 - x |
---|
209 | dx = x2 - x1 |
---|
210 | dx1 = dx1/dx |
---|
211 | dx2 = dx2/dx |
---|
212 | |
---|
213 | y1 = yp[{zdim:idk1}] |
---|
214 | y2 = yp[{zdim:idk2}] |
---|
215 | |
---|
216 | if 'linear' in method : |
---|
217 | result = dx1*y2 + dx2*y1 |
---|
218 | if 'log' in method : |
---|
219 | if yp_min > 0. : |
---|
220 | result = np.power(y2, dx1) * np.power(y1, dx2) |
---|
221 | if yp_max < 0. : |
---|
222 | result = -np.power(-y2, dx1) * np.power(-y1, dx2) |
---|
223 | if 'nearest' in method : |
---|
224 | result = xr.where ( dx2>=dx1, y1, y2) |
---|
225 | |
---|
226 | return result |
---|
227 | |
---|
228 | for k in np.arange (nk_ou) : |
---|
229 | result = __interp (x[{pdim:k}], xp, yp) |
---|
230 | |
---|
231 | # Put result in the final array |
---|
232 | ou_tab [{pdim:k}] = result |
---|
233 | |
---|
234 | return ou_tab.squeeze() |
---|
235 | |
---|
236 | def fixed_lon (lon, center_lon=0.0) : |
---|
237 | '''Returns corrected longitudes for nicer plots |
---|
238 | |
---|
239 | lon : longitudes of the grid. At least 1D. |
---|
240 | center_lon : center longitude. Default=0. |
---|
241 | |
---|
242 | Designed by Phil Pelson. See https://gist.github.com/pelson/79cf31ef324774c97ae7 |
---|
243 | ''' |
---|
244 | mmath = __mmath__ (lon) |
---|
245 | |
---|
246 | zfixed_lon = lon.copy () |
---|
247 | |
---|
248 | zfixed_lon = mmath.where (zfixed_lon > center_lon+180., zfixed_lon-360.0, fixed_lon) |
---|
249 | zfixed_lon = mmath.where (zfixed_lon < center_lon-180., zfixed_lon+360.0, fixed_lon) |
---|
250 | |
---|
251 | start = np.argmax (np.abs (np.diff (zfixed_lon, axis=-1)) > 180., axis=-1) |
---|
252 | zfixed_lon [start+1:] += 360. |
---|
253 | |
---|
254 | return zfixed_lon |
---|
255 | |
---|
256 | def nord2sud (p2d) : |
---|
257 | '''Swap north to south a 2D field |
---|
258 | ''' |
---|
259 | pout = p2d [..., -1::-1, : ] |
---|
260 | |
---|
261 | return pout |
---|
262 | |
---|
263 | def point2geo (p1d, verbose=False, lon=False, lat=False, jpi=0, jpj=0, share_pole=False) : |
---|
264 | ''' |
---|
265 | From 1D (..., points_physiques) (restart type) to 2D (..., lat, lon) |
---|
266 | ''' |
---|
267 | math = __mmath__ (p1d) |
---|
268 | |
---|
269 | # Get the horizontal dimension |
---|
270 | jpn = p1d.shape[-1] |
---|
271 | # Get other dimension |
---|
272 | form1 = list(p1d.shape [0:-1]) |
---|
273 | |
---|
274 | if jpi==0 and jpj>0 : |
---|
275 | jpi = (jpn-2)//(jpj-2) |
---|
276 | if jpi>0 and jpj == 0 : |
---|
277 | jpj = (jpn-2)//jpi +2 |
---|
278 | |
---|
279 | def c_jpn (jpj, jpi) : return jpi*(jpj-2) + 2 |
---|
280 | |
---|
281 | if jpi==0 and jpi==0 : |
---|
282 | if jpn == c_jpn( 95, 96) : |
---|
283 | jpj, jpi = 95, 96 |
---|
284 | if jpn == c_jpn (96, 96) : |
---|
285 | jpj, jpi = 96, 96 |
---|
286 | if jpn == c_jpn (144, 144) : |
---|
287 | jpj, jpi = 144, 144 |
---|
288 | if jpn == c_jpn (143, 144) : |
---|
289 | jpj, jpi = 143, 144 |
---|
290 | |
---|
291 | form2 = form1 + [jpj, jpi,] |
---|
292 | form_shape = form1 + [jpj-2, jpi] |
---|
293 | p2d = np.empty ( form2 ) |
---|
294 | |
---|
295 | p2d [..., 1:jpj-1, :] = np.reshape (np.float64 (p1d [..., 1:jpn-1]), form_shape ) |
---|
296 | if share_pole : |
---|
297 | for ji in np.arange ( p2d.shape[-1] ) : |
---|
298 | p2d [..., 0 , ji] = p1d[..., 0 ] / float(jpi) |
---|
299 | p2d [..., jpj-1 , ji] = p1d[..., jpn-1] / float(jpi) |
---|
300 | else : |
---|
301 | for ji in np.arange ( p2d.shape[-1] ) : |
---|
302 | p2d [..., 0 , ji] = p1d[..., 0 ] |
---|
303 | p2d [..., jpj-1 , ji] = p1d[..., jpn-1] |
---|
304 | |
---|
305 | if math == xr : |
---|
306 | p2d = xr.DataArray (p2d) |
---|
307 | p2d.attrs.update ( p1d.attrs ) |
---|
308 | for idim in np.arange ( len(p1d.shape [0:-1]) ): |
---|
309 | dim = p1d.dims[idim] |
---|
310 | p2d = p2d.rename ( {p2d.dims[idim]:p1d.dims[idim]} ) |
---|
311 | p2d = p2d.assign_coords ( {p2d.dims[idim]:p1d.coords[dim]} ) |
---|
312 | p2d = p2d.rename ( {p2d.dims[-1]:'x', p2d.dims[-2]:'y'} ) |
---|
313 | p2d['x'].attrs.update ( {'axis':'X'} ) |
---|
314 | p2d['y'].attrs.update ( {'axis':'Y'} ) |
---|
315 | |
---|
316 | if type(lon) == bool : |
---|
317 | if lon : |
---|
318 | lon = np.linspace ( -180, 180, jpi, endpoint=False) |
---|
319 | #if math == xr : |
---|
320 | # lon = xr.DataArray ( lon, dims=('lon',), coords=(lon,) ) |
---|
321 | if type(lat) == bool : |
---|
322 | if lat : |
---|
323 | lat = np.linspace ( 90, -90, jpj, endpoint=True) |
---|
324 | #if math == xr : |
---|
325 | # lat = xr.DataArray ( lat, dims=('lat',), coords=(lat,) ) |
---|
326 | if type(lon) != bool and type(lat) != bool : |
---|
327 | if __mmath__ (lat) == xr : |
---|
328 | lat_name = lat.name |
---|
329 | else : |
---|
330 | lat_name = 'lat' |
---|
331 | if __mmath__ (lon) == xr : |
---|
332 | lon_name = lon.name |
---|
333 | else : |
---|
334 | lon_name = 'lon' |
---|
335 | p2d = p2d.rename ( {p2d.dims[-1]:lon_name, p2d.dims[-2]:lat_name} ) |
---|
336 | p2d = p2d.assign_coords ( {lon_name:lon, lat_name:lat} ) |
---|
337 | p2d[lon_name].attrs.update ( { 'units':'degrees_east' , 'long_name':'Longitude', 'standard_name':'longitude', 'axis':'X' } ) |
---|
338 | p2d[lat_name].attrs.update ( { 'units':'degrees_north', 'long_name':'Latitude' , 'standard_name':'latitude' , 'axis':'Y' } ) |
---|
339 | |
---|
340 | return p2d |
---|
341 | |
---|
342 | def geo2point ( p2d, cumul_poles=False, dim1d='points_physiques' ) : |
---|
343 | ''' |
---|
344 | From 2D (..., lat, lon) to 1D (..., points_phyiques) |
---|
345 | ''' |
---|
346 | math = __mmath__ (p2d) |
---|
347 | # |
---|
348 | # Get the horizontal dimensions |
---|
349 | (jpj, jpi) = p2d.shape[-2:] |
---|
350 | jpn = jpi*(jpj-2) + 2 |
---|
351 | |
---|
352 | # Get other dimensions |
---|
353 | form1 = list(p2d.shape [0:-2]) |
---|
354 | form2 = form1 + [jpn,] |
---|
355 | |
---|
356 | p1d = np.empty ( form2 ) |
---|
357 | form_shape = form1 + [jpn-2,] |
---|
358 | |
---|
359 | if math == xr : |
---|
360 | p1d[..., 1:-1] = np.reshape ( np.float64 (p2d[..., 1:-1, :].values).ravel(), form_shape ) |
---|
361 | else : |
---|
362 | p1d[..., 1:-1] = np.reshape ( np.float64 (p2d[..., 1:-1, :] ).ravel(), form_shape ) |
---|
363 | p1d[..., 0 ] = p2d[..., 0, 0] |
---|
364 | p1d[..., -1 ] = p2d[..., -1, 0] |
---|
365 | |
---|
366 | if math == xr : |
---|
367 | p1d = xr.DataArray (p1d) |
---|
368 | p1d.attrs.update ( p2d.attrs ) |
---|
369 | if len(p2d.shape [0:-2]) > 0 : |
---|
370 | for idim in np.arange ( len(p2d.shape [0:-2]) ): |
---|
371 | dim=p2d.dims[idim] |
---|
372 | p1d = p1d.rename ( {p1d.dims[idim]:p2d.dims[idim]} ) |
---|
373 | p1d = p1d.assign_coords ( {p1d.dims[idim] :p2d.coords[dim].values} ) |
---|
374 | p1d = p1d.rename ( {p1d.dims[-1]:dim1d} ) |
---|
375 | |
---|
376 | if cumul_poles : |
---|
377 | p1d[..., 0] = np.sum ( p2d[..., 0, :] ) |
---|
378 | p1d[..., -1] = np.sum ( p2d[..., -1, :] ) |
---|
379 | |
---|
380 | return p1d |
---|
381 | |
---|
382 | def geo2en (pxx, pyy, pzz, glam, gphi) : |
---|
383 | '''Change vector from geocentric to east/north |
---|
384 | |
---|
385 | Inputs : |
---|
386 | pxx, pyy, pzz : components on the geocentric system |
---|
387 | glam, gphi : longitude and latitude of the points |
---|
388 | ''' |
---|
389 | |
---|
390 | gsinlon = np.sin (RAD * glam) |
---|
391 | gcoslon = np.cos (RAD * glam) |
---|
392 | gsinlat = np.sin (RAD * gphi) |
---|
393 | gcoslat = np.cos (RAD * gphi) |
---|
394 | |
---|
395 | pte = - pxx * gsinlon + pyy * gcoslon |
---|
396 | ptn = - pxx * gcoslon * gsinlat - pyy * gsinlon * gsinlat + pzz * gcoslat |
---|
397 | |
---|
398 | return pte, ptn |
---|
399 | |
---|
400 | def en2geo (pte, ptn, glam, gphi) : |
---|
401 | '''Change vector from east/north to geocentric |
---|
402 | |
---|
403 | Inputs : |
---|
404 | pte, ptn : eastward/northward components |
---|
405 | glam, gphi : longitude and latitude of the points |
---|
406 | ''' |
---|
407 | |
---|
408 | gsinlon = np.sin (RAD * glam) |
---|
409 | gcoslon = np.cos (RAD * glam) |
---|
410 | gsinlat = np.sin (RAD * gphi) |
---|
411 | gcoslat = np.cos (RAD * gphi) |
---|
412 | |
---|
413 | pxx = - pte * gsinlon - ptn * gcoslon * gsinlat |
---|
414 | pyy = pte * gcoslon - ptn * gsinlon * gsinlat |
---|
415 | pzz = ptn * gcoslat |
---|
416 | |
---|
417 | return pxx, pyy, pzz |
---|
418 | |
---|
419 | ## =========================================================================== |
---|
420 | ## |
---|
421 | ## That's all folk's !!! |
---|
422 | ## |
---|
423 | ## =========================================================================== |
---|