[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$" |
---|
| 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 __math__ (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 = __math__ (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', 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 ( 'in_shape : ', in_shape ) |
---|
| 131 | ou_shape = np.array (in_shape) |
---|
| 132 | if verbose : print ( 'ou_shape : ', ou_shape ) |
---|
| 133 | ou_shape[axis] = nk_ou |
---|
| 134 | |
---|
| 135 | in_dims = list (yp.dims) |
---|
| 136 | if verbose : print ( 'in_dims : ', in_dims ) |
---|
| 137 | ou_dims = in_dims |
---|
| 138 | if verbose : print ( 'ou_dims : ', ou_dims ) |
---|
| 139 | pdim = x.dims[0] |
---|
| 140 | ou_dims[axis] = pdim |
---|
| 141 | |
---|
| 142 | new_coords = [] |
---|
| 143 | for coord in yp.dims : |
---|
| 144 | if coord == zdim : new_coords.append (x.coords[pdim].values) |
---|
| 145 | else : new_coords.append (yp.coords[coord].values) |
---|
| 146 | |
---|
| 147 | if verbose : print ( 'new_coords : ', new_coords ) |
---|
| 148 | |
---|
| 149 | ou_tab = xr.DataArray (np.empty (ou_shape), dims=ou_dims, coords=new_coords) |
---|
| 150 | |
---|
| 151 | if 'log' in method : |
---|
| 152 | yp_min = yp.min() ; yp_max = yp.max() |
---|
| 153 | indic = yp_min * yp_max |
---|
| 154 | if indic <= 0. : |
---|
| 155 | print ('Input data have a change of sign') |
---|
| 156 | print ('Error : logarithmic method is available only for') |
---|
| 157 | print ('positive or negative input values. ') |
---|
| 158 | raise Exception |
---|
| 159 | |
---|
| 160 | # Optimized (pre-compiled) interpolation loop |
---|
| 161 | #@numba.jit(nopython=True) |
---|
| 162 | def __interp (nk_ou, x, xp, yp, result) : |
---|
| 163 | # Interpolate |
---|
| 164 | for k in np.arange (nk_ou) : |
---|
| 165 | # Find index of the just above level |
---|
| 166 | idk1 = np.minimum ( (x[k]-xp), 0.).argmax (dim=zdim) |
---|
| 167 | idk2 = idk1 - 1 |
---|
| 168 | idk2 = np.maximum (idk2, 0) |
---|
| 169 | |
---|
| 170 | x1 = xp[{zdim:idk1}] |
---|
| 171 | x2 = xp[{zdim:idk2}] |
---|
| 172 | |
---|
| 173 | dx1 = x[k] - x1 |
---|
| 174 | dx2 = x2 - x[k] |
---|
| 175 | dx = x2 - x1 |
---|
| 176 | dx1 = dx1/dx ; dx2 = dx2/dx |
---|
| 177 | |
---|
| 178 | y1 = yp[{zdim:idk1}] |
---|
| 179 | y2 = yp[{zdim:idk2}] |
---|
| 180 | |
---|
| 181 | if 'linear' in method : |
---|
| 182 | result = (dx1*y2 + dx2*y1) |
---|
| 183 | if 'log' in method : |
---|
| 184 | if yp_min > 0. : |
---|
| 185 | result = np.power(y2, dx1) * np.power(y1, dx2) |
---|
| 186 | if yp_max < 0. : |
---|
| 187 | result = -np.power(-y2, dx1) * np.power(-y1, dx2) |
---|
| 188 | if 'nearest' in method : |
---|
| 189 | result = xr.where ( dx2>=dx1, y1, y2) |
---|
| 190 | |
---|
| 191 | return result |
---|
| 192 | |
---|
| 193 | result = __interp (nk_ou, x, xp, yp, result) |
---|
| 194 | |
---|
| 195 | # Put result in the final array |
---|
| 196 | ou_tab [{pdim:k}] = result |
---|
| 197 | |
---|
| 198 | return ou_tab.squeeze() |
---|
| 199 | |
---|
| 200 | def nord2sud (p2D) : |
---|
| 201 | ''' |
---|
| 202 | Swap north to south a 2D field |
---|
| 203 | ''' |
---|
| 204 | pout = p2D [..., -1::-1, : ] |
---|
| 205 | |
---|
| 206 | return pout |
---|
| 207 | |
---|
| 208 | def point2geo (p1D) : |
---|
| 209 | ''' |
---|
| 210 | From 1D (restart type) to 2D |
---|
| 211 | ''' |
---|
| 212 | math = __math__ (p1D) |
---|
| 213 | |
---|
| 214 | # Get the dimensions |
---|
| 215 | jpn = p1D.shape[-1] |
---|
| 216 | |
---|
| 217 | if len (p1D.shape) > 1 : |
---|
| 218 | jpt = p1D.shape[0] |
---|
| 219 | else : |
---|
| 220 | jpt = 0 |
---|
| 221 | |
---|
| 222 | if jpn == 9026 : jpi = 96 ; jpj = 96 |
---|
| 223 | if jpn == 17858 : jpi = 144 ; jpj = 144 |
---|
| 224 | if jpn == 20306 : jpi = 144 ; jpj = 143 |
---|
| 225 | |
---|
| 226 | if jpt > 0 : |
---|
| 227 | p2D = np.zeros ( (jpt, jpj, jpi) ) |
---|
| 228 | p2D [:, 1:jpj-1, :] = np.reshape (p1D [:,1:jpn-1], (jpt, jpj-2, jpi) ) |
---|
| 229 | p2D [:, 0 , : ] = p1D[:, 0 ] |
---|
| 230 | p2D [:, jpj-1, : ] = p1D[:, jpn-1] |
---|
| 231 | |
---|
| 232 | else : |
---|
| 233 | p2D = np.zeros ( (jpj, jpi) ) |
---|
| 234 | p2D [1:jpj-1, :] = np.reshape (np.float64 (p1D [1:jpn-1]), (jpj-2, jpi) ) |
---|
| 235 | p2D [0 , : ] = p1D[ 0 ] |
---|
| 236 | p2D [jpj-1, : ] = p1D[jpn-1] |
---|
| 237 | |
---|
| 238 | if math == xr : |
---|
| 239 | p2D = xr.DataArray (p2D) |
---|
[6508] | 240 | for attr in p1D.attrs : |
---|
| 241 | p2D.attrs[attr] = p1D.attrs[attr] |
---|
[6265] | 242 | p2D = p2D.rename ( {p2D.dims[0]:p1D.dims[0], p2D.dims[-1]:'x', p2D.dims[-2]:'y'} ) |
---|
| 243 | |
---|
| 244 | return p2D |
---|
| 245 | |
---|
[6508] | 246 | def geo2point ( p2D, cumulPoles=False, dim1D='cell', debug=False ) : |
---|
[6265] | 247 | ''' |
---|
| 248 | From 2D (lat, lon) to 1D (points_phyiques) |
---|
| 249 | ''' |
---|
| 250 | math = __math__ (p2D) |
---|
| 251 | # |
---|
| 252 | # Get the dimension |
---|
| 253 | |
---|
| 254 | (jpj, jpi) = p2D.shape[-2:] |
---|
| 255 | |
---|
| 256 | if len (p2D.shape) > 2 : |
---|
| 257 | jpt = p2D.shape[0] |
---|
| 258 | else : |
---|
| 259 | jpt = -1 |
---|
| 260 | |
---|
| 261 | jpn = jpi*(jpj-2) + 2 |
---|
| 262 | |
---|
[6508] | 263 | if debug : |
---|
| 264 | print ( f'{jpi=} {jpj=} {jpn=} {jpt=}' ) |
---|
| 265 | |
---|
[6265] | 266 | if jpt == -1 : |
---|
| 267 | p1D = np.zeros ( (jpn,) ) |
---|
| 268 | p1D[1:-1] = np.float64(p2D[1:-1, :]).ravel() |
---|
| 269 | p1D[ 0] = p2D[ 0, 0] |
---|
| 270 | p1D[-1] = p2D[-1, 0] |
---|
| 271 | |
---|
| 272 | else : |
---|
| 273 | p1D = np.zeros ( (jpt, jpn) ) |
---|
| 274 | if math == xr : |
---|
| 275 | p1D[:, 1:-1] = np.reshape ( np.float64 (p2D[:, 1:-1, :].values).ravel(), (jpt, jpn-2) ) |
---|
| 276 | else : |
---|
| 277 | p1D[:, 1:-1] = np.reshape ( np.float64 (p2D[:, 1:-1, :] ).ravel(), (jpt, jpn-2) ) |
---|
| 278 | p1D[:, 0 ] = p2D[:, 0, 0] |
---|
| 279 | p1D[:, -1 ] = p2D[:, -1, 0] |
---|
| 280 | |
---|
| 281 | if math == xr : |
---|
| 282 | p1D = xr.DataArray (p1D) |
---|
[6508] | 283 | for attr in p2D.attrs : |
---|
| 284 | p1D.attrs[attr] = p2D.attrs[attr] |
---|
[6265] | 285 | p1D = p1D.rename ( {p1D.dims[0]:p2D.dims[0], p1D.dims[-1]:dim1D} ) |
---|
| 286 | |
---|
| 287 | if cumulPoles : |
---|
| 288 | p1D[..., 0] = np.sum ( p2D[..., 0, :] ) |
---|
| 289 | p1D[..., -1] = np.sum ( p2D[..., -1, :] ) |
---|
| 290 | |
---|
| 291 | return p1D |
---|
| 292 | |
---|
[6508] | 293 | def geo3point ( p3D, cumulPoles=False, dim1D='cell' ) : |
---|
[6265] | 294 | ''' |
---|
| 295 | From 3D (lev, lat, lon) to 2D (lev, points_phyiques) |
---|
| 296 | ''' |
---|
| 297 | math = __math__ (p3D) |
---|
| 298 | # |
---|
| 299 | # Get the dimensions |
---|
| 300 | |
---|
| 301 | (jpk, jpj, jpi) = p3D.shape[-3:] |
---|
| 302 | |
---|
| 303 | if len (p3D.shape) > 3 : |
---|
| 304 | jpt = p3D.shape[0] |
---|
| 305 | else : |
---|
| 306 | jpt = -1 |
---|
| 307 | |
---|
| 308 | jpn = jpi*(jpj-2) + 2 |
---|
| 309 | |
---|
| 310 | if jpt == -1 : |
---|
| 311 | p2D = np.zeros ( (jpk, jpn,) ) |
---|
| 312 | for jk in np.arange (jpk) : |
---|
| 313 | p2D [jk, :] = geo2point ( p3D [jk,:,:], cumulPoles, dim1D ) |
---|
| 314 | else : |
---|
| 315 | p2D = np.zeros ( (jpt, jpk, jpn) ) |
---|
| 316 | for jk in np.arange (jpk) : |
---|
| 317 | p2D [:, jk, :] = geo2point ( p3D [:, jk,:,:], cumulPoles, dim1D ) |
---|
| 318 | |
---|
| 319 | if math == xr : |
---|
| 320 | p2D = xr.DataArray (p2D) |
---|
[6508] | 321 | for attr in p2D.attrs : |
---|
| 322 | p2D.attrs[attr] = p3D.attrs[attr] |
---|
[6265] | 323 | p2D = p2D.rename ( {p2D.dims[-1]:dim1D, p2D.dims[-2]:p3D.dims[-3]} ) |
---|
| 324 | |
---|
| 325 | return p2D |
---|
| 326 | |
---|
| 327 | def geo2en (pxx, pyy, pzz, glam, gphi) : |
---|
| 328 | ''' |
---|
| 329 | Change vector from geocentric to east/north |
---|
| 330 | |
---|
| 331 | Inputs : |
---|
| 332 | pxx, pyy, pzz : components on the geocentric system |
---|
| 333 | glam, gphi : longitude and latitude of the points |
---|
| 334 | ''' |
---|
| 335 | |
---|
| 336 | gsinlon = np.sin (rad * glam) |
---|
| 337 | gcoslon = np.cos (rad * glam) |
---|
| 338 | gsinlat = np.sin (rad * gphi) |
---|
| 339 | gcoslat = np.cos (rad * gphi) |
---|
| 340 | |
---|
| 341 | pte = - pxx * gsinlon + pyy * gcoslon |
---|
| 342 | ptn = - pxx * gcoslon * gsinlat - pyy * gsinlon * gsinlat + pzz * gcoslat |
---|
| 343 | |
---|
| 344 | return pte, ptn |
---|
| 345 | |
---|
| 346 | def en2geo (pte, ptn, glam, gphi) : |
---|
| 347 | ''' |
---|
| 348 | Change vector from east/north to geocentric |
---|
| 349 | |
---|
| 350 | Inputs : |
---|
| 351 | pte, ptn : eastward/northward components |
---|
| 352 | glam, gphi : longitude and latitude of the points |
---|
| 353 | ''' |
---|
| 354 | |
---|
| 355 | gsinlon = np.sin (rad * glam) |
---|
| 356 | gcoslon = np.cos (rad * glam) |
---|
| 357 | gsinlat = np.sin (rad * gphi) |
---|
| 358 | gcoslat = np.cos (rad * gphi) |
---|
| 359 | |
---|
| 360 | pxx = - pte * gsinlon - ptn * gcoslon * gsinlat |
---|
| 361 | pyy = pte * gcoslon - ptn * gsinlon * gsinlat |
---|
| 362 | pzz = ptn * gcoslat |
---|
| 363 | |
---|
| 364 | return pxx, pyy, pzz |
---|
| 365 | |
---|
| 366 | ## =========================================================================== |
---|
| 367 | ## |
---|
| 368 | ## That's all folk's !!! |
---|
| 369 | ## |
---|
| 370 | ## =========================================================================== |
---|