[6265] | 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$" |
---|
[6647] | 30 | __Id__ = "$Id: $" |
---|
[6265] | 31 | __HeadURL = "$HeadURL$" |
---|
| 32 | |
---|
| 33 | import numpy as np |
---|
| 34 | |
---|
| 35 | try : import xarray as xr |
---|
| 36 | except ImportError : pass |
---|
| 37 | |
---|
[6647] | 38 | try : import numba |
---|
| 39 | except ImportError : pass |
---|
[6265] | 40 | |
---|
| 41 | rpi = np.pi ; rad = np.deg2rad (1.0) ; dar = np.rad2deg (1.0) |
---|
| 42 | |
---|
| 43 | |
---|
[6647] | 44 | def __mmath__ (tab) : |
---|
[6265] | 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 | ''' |
---|
[6647] | 73 | math = __mmath__ (tab) |
---|
[6265] | 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 | |
---|
[6647] | 101 | def interp1d (x, xp, yp, zdim='presnivs', units=None, verbose=False, method='linear') : |
---|
[6265] | 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 \!/ |
---|
[6647] | 120 | ''' |
---|
[6265] | 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) |
---|
[6647] | 130 | if verbose : print ( f'{in_shape=}' ) |
---|
[6265] | 131 | ou_shape = np.array (in_shape) |
---|
[6647] | 132 | if verbose : print ( f'{ou_shape=}' ) |
---|
[6265] | 133 | ou_shape[axis] = nk_ou |
---|
| 134 | |
---|
| 135 | in_dims = list (yp.dims) |
---|
[6647] | 136 | if verbose : print ( f'{in_dims=}' ) |
---|
[6265] | 137 | ou_dims = in_dims |
---|
[6647] | 138 | |
---|
[6265] | 139 | pdim = x.dims[0] |
---|
| 140 | ou_dims[axis] = pdim |
---|
| 141 | |
---|
| 142 | new_coords = [] |
---|
[6647] | 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) |
---|
[6265] | 150 | |
---|
[6647] | 151 | if verbose : |
---|
| 152 | print ( f'{ou_dims =}' ) |
---|
| 153 | print ( f'{new_coords=}' ) |
---|
[6265] | 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 |
---|
[6647] | 167 | #@numba.jit (nopython=True) |
---|
| 168 | def __interp (nk_ou, x, xp, yp) : |
---|
[6265] | 169 | # Interpolate |
---|
[6647] | 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) |
---|
[6265] | 197 | |
---|
[6647] | 198 | return result |
---|
[6265] | 199 | |
---|
[6647] | 200 | for k in np.arange (nk_ou) : |
---|
| 201 | result = __interp (nk_ou, x[{pdim:k}], xp, yp) |
---|
| 202 | |
---|
[6265] | 203 | # Put result in the final array |
---|
| 204 | ou_tab [{pdim:k}] = result |
---|
| 205 | |
---|
| 206 | return ou_tab.squeeze() |
---|
| 207 | |
---|
[6647] | 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 | |
---|
[6265] | 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 | ''' |
---|
[6647] | 241 | math = __mmath__ (p1D) |
---|
[6265] | 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) |
---|
[6508] | 269 | for attr in p1D.attrs : |
---|
| 270 | p2D.attrs[attr] = p1D.attrs[attr] |
---|
[6265] | 271 | p2D = p2D.rename ( {p2D.dims[0]:p1D.dims[0], p2D.dims[-1]:'x', p2D.dims[-2]:'y'} ) |
---|
| 272 | |
---|
| 273 | return p2D |
---|
| 274 | |
---|
[6647] | 275 | def geo2point ( p2D, cumulPoles=False, dim1D='points_physiques' ) : |
---|
[6265] | 276 | ''' |
---|
| 277 | From 2D (lat, lon) to 1D (points_phyiques) |
---|
| 278 | ''' |
---|
[6647] | 279 | math = __mmath__ (p2D) |
---|
[6265] | 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) |
---|
[6647] | 309 | for attr in p2D.attrs : |
---|
[6508] | 310 | p1D.attrs[attr] = p2D.attrs[attr] |
---|
[6265] | 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 | |
---|
[6647] | 319 | def geo3point ( p3D, cumulPoles=False, dim1D='points_physiques' ) : |
---|
[6265] | 320 | ''' |
---|
| 321 | From 3D (lev, lat, lon) to 2D (lev, points_phyiques) |
---|
| 322 | ''' |
---|
[6647] | 323 | math = __mmath__ (p3D) |
---|
[6265] | 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) |
---|
[6508] | 347 | for attr in p2D.attrs : |
---|
| 348 | p2D.attrs[attr] = p3D.attrs[attr] |
---|
[6265] | 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 | ## =========================================================================== |
---|