Changeset 6064 for TOOLS/MOSAIX/nemo.py
- Timestamp:
- 02/24/22 09:58:49 (2 years ago)
- File:
-
- 1 edited
Legend:
- Unmodified
- Added
- Removed
-
TOOLS/MOSAIX/nemo.py
r5948 r6064 18 18 ''' 19 19 20 import sys 21 22 try : import numpy as np 23 except : pass 24 20 import sys, numpy as np 25 21 try : import xarray as xr 26 22 except : pass 23 24 rpi = np.pi ; rad = rpi / 180.0 25 26 nperio_valid_range = [0,1,4,5,6] 27 27 28 28 ## SVN information … … 33 33 __HeadURL = "$HeadURL$" 34 34 35 def __guessNperio__ (jp i, nperio=None) :35 def __guessNperio__ (jpj, jpi, nperio=None) : 36 36 ''' 37 37 Tries to guess the value of nperio (periodicity parameter. See NEMO documentation for details) … … 41 41 nperio : periodicity parameter 42 42 ''' 43 44 43 if nperio == None : 45 44 if jpi == 182 : nperio = 4 # ORCA2. We choose legacy orca2. 46 if jpi == 362 : nperio = 6 # ORCA1. 47 if jpi == 1442 : nperio = 6 # ORCA025. 45 if jpi == 362 : nperio = 6 # eORCA1. 46 if jpi == 1442 : nperio = 6 # ORCA025. 47 48 if jpj == 149 : nperio = 4 # ORCA2. We choose legacy orca2. 49 if jpj == 332 : nperio = 6 # eORCA1. 48 50 # 49 51 if nperio == None : 50 sys.exit ('in nemo .lbc: nperio not found, and cannot by guessed' )52 sys.exit ('in nemo module : nperio not found, and cannot by guessed' ) 51 53 else : 52 print ('nperio set as {:d} (deduced from jpi={:d})'.format (nperio, jpi))54 print ('nperio set as {:d} (deduced from jpi={:d})'.format (nperio, jpi)) 53 55 54 56 return nperio 57 58 def __guessPoint__ (ptab) : 59 ''' 60 Tries to guess the grid point (periodicity parameter. See NEMO documentation for details) 61 62 For array conforments with xgcm requirements 63 64 Inputs 65 ptab : xarray array 66 67 Credits : who is the original author ? 68 ''' 69 gP = None 70 if isinstance (ptab, xr.core.dataarray.DataArray) : 71 if 'x_c' in ptab.dims and 'y_c' in ptab.dims : gP = 'T' 72 if 'x_f' in ptab.dims and 'y_c' in ptab.dims : gP = 'U' 73 if 'x_c' in ptab.dims and 'y_f' in ptab.dims : gP = 'V' 74 if 'x_f' in ptab.dims and 'y_f' in ptab.dims : gP = 'F' 75 if 'x_c' in ptab.dims and 'y_c' in ptab.dims and 'z_c' in ptab.dims : gP = 'T' 76 if 'x_c' in ptab.dims and 'y_c' in ptab.dims and 'z_f' in ptab.dims : gP = 'W' 77 if 'x_f' in ptab.dims and 'y_c' in ptab.dims and 'z_f' in ptab.dims : gP = 'U' 78 if 'x_c' in ptab.dims and 'y_f' in ptab.dims and 'z_f' in ptab.dims : gP = 'V' 79 if 'x_f' in ptab.dims and 'y_f' in ptab.dims and 'z_f' in ptab.dims : gP = 'F' 80 81 if gP == None : 82 sys.exit ('in nemo module : cd_type not found, and cannot by guessed' ) 83 else : 84 print ('Grid set as', gP, 'deduced from dims ', ptab.dims) 85 return gP 86 else : 87 sys.exit ('in nemo module : cd_type not found, input is not an xarray data' ) 88 #return gP 55 89 56 90 def fixed_lon (lon) : … … 60 94 fixed_lon (lon) 61 95 lon : longitudes of the grid. At least 2D. 96 97 Designed by Phil Pelson. See https://gist.github.com/pelson/79cf31ef324774c97ae7 62 98 ''' 63 99 fixed_lon = lon.copy () 64 100 for i, start in enumerate (np.argmax (np.abs (np.diff (lon, axis=-1)) > 180, axis=-1)) : 65 fixed_lon[...,i, start+1:] += 360 101 fixed_lon[..., i, start+1:] += 360 102 103 # Special case for eORCA025 104 if lon.shape[-1] == 1442 : lon [..., -2, :] = lon [..., -3, :] 105 66 106 return fixed_lon 67 107 … … 82 122 lat (optionnal) : latitudes of the grid 83 123 ''' 84 85 124 if np.max (lat) != None : 86 125 je = jeq (lat) … … 95 134 return lon1D 96 135 97 def latreg (lat) : 98 ''' 99 Returns maximum j index wehre gridlines are along latitude in the norethern hemisphere 100 101 lon : longitudes of the grid 102 lat (optionnal) : latitudes of the grid 103 ''' 104 dy = np.float64 (np.mean (np.abs (lat - np.roll(lat,shift=1,axis=-2)))) 136 def latreg (lat, diff=0.1) : 137 ''' 138 Returns maximum j index where gridlines are along latitude in the northern hemisphere 139 140 lat : latitudes of the grid (2D) 141 diff [optional] : tolerance 142 ''' 143 if diff == None : 144 dy = np.float64 (np.mean (np.abs (lat - np.roll(lat,shift=1,axis=-2)))) 145 diff = dy/100. 146 105 147 je = jeq (lat) 106 jreg = np.where (lat[...,je:,:].max(axis=-1) - lat[...,je:,:].min(axis=-1)< d y/100.)[-1][-1] + je148 jreg = np.where (lat[...,je:,:].max(axis=-1) - lat[...,je:,:].min(axis=-1)< diff)[-1][-1] + je 107 149 latreg = np.float64 (lat[...,jreg,:].mean(axis=-1)) 108 JREG = jreg150 JREG = jreg 109 151 110 152 return jreg, latreg … … 114 156 Returns 1D latitudes for zonal means and simple plots. 115 157 116 lat : latitudes of the grid 158 lat : latitudes of the grid (2D) 117 159 ''' 118 160 jpj, jpi = lat.shape[-2:] 119 161 try : 120 if type (lat) == 162 if type (lat) == xr.core.dataarray.DataArray : math = xr 121 163 else : math = np 122 164 except : math = np 123 165 124 # jreg : index of last uniform latitude, north of equator 125 126 dy = np.float64 (np.mean (np.abs (lat - np.roll (lat, shift=1,axis=-2)))) 166 dy = np.float64 (np.mean (np.abs (lat - np.roll (lat, shift=1,axis=-2)))) 127 167 je = jeq (lat) 128 lat eq =np.float64 (lat[...,je,:].mean(axis=-1))168 lat_eq = np.float64 (lat[...,je,:].mean(axis=-1)) 129 169 130 jreg = np.where (lat[...,je:,:].max(axis=-1) - lat[...,je:,:].min(axis=-1)< dy/100.)[-1][-1] + je 131 latreg = np.float64 (lat[...,jreg,:].mean(axis=-1)) 132 latave = np.mean (lat, axis=-1) 133 134 if ( np.abs (lateq) < dy/100. ) : # T, U or W grid 135 dys = (90.-latreg)//(jpj-jreg-1)*0.5 136 yrange = 90.-dys-latreg 137 else : # V or F grid 138 yrange = 90. -latreg 139 140 lat1D = math.where (latave<latreg, latave, latreg + yrange * (np.arange(jpj)-jreg)/(jpj-jreg-1) ) 141 142 if math == xr : 143 lat1D.attrs = lat.attrs 170 jreg, lat_reg = latreg (lat) 171 lat_ave = np.mean (lat, axis=-1) 172 173 if ( np.abs (lat_eq) < dy/100. ) : # T, U or W grid 174 dys = (90.-lat_reg) / (jpj-jreg-1)*0.5 175 yrange = 90.-dys-lat_reg 176 else : # V or F grid 177 yrange = 90. -lat_reg 178 179 lat1D = math.where (lat_ave<lat_reg, lat_ave, lat_reg + yrange * (np.arange(jpj)-jreg)/(jpj-jreg-1) ) 180 181 if math == xr : lat1D.attrs = lat.attrs 144 182 145 183 return lat1D … … 148 186 ''' 149 187 Returns simple latitude and longitude (1D) for simple plots. 150 ''' 151 laty = lat1D (lat) 152 lonx = lon1D (lon, lat) 153 154 return laty, lonx 155 156 def extend (tab, Lon=False, jplus=25) : 157 ''' 158 Returns extended field eastward to have better plots 188 189 lat, lon : latitudes and longitudes of the grid (2D) 190 ''' 191 return lat1D (lat), lon1D (lon, lat) 192 193 def mask_lonlat (ptab, x0, x1, y0, y1, lon, lat, sval=np.nan) : 194 try : 195 lon = lon.to_masked_array() 196 lat = lat.to_masked_array() 197 except : pass 198 199 mask = np.logical_and ( np.logical_and(lat>y0, lat<y1), 200 np.logical_or (np.logical_or (np.logical_and(lon>x0, lon<x1), np.logical_and(lon+360>x0, lon+360<x1)), 201 np.logical_and(lon-360>x0, lon-360<x1)) ) 202 try : tab = xr.where (mask, ptab, np.nan) 203 except : tab = np.where (mask, ptab, np.nan) 204 205 return tab 206 207 def extend (tab, Lon=False, jplus=25, jpi=None, nperio=4) : 208 ''' 209 Returns extended field eastward to have better plots, and box average crossing the boundary 159 210 Works only for xarray and numpy data (?) 160 211 161 212 tab : field to extend. 162 213 Lon : (optional, default=False) : if True, add 360 in the extended parts of the field 214 jpi : normal longitude dimension of the field. exrtend does nothing it the actual 215 size of the field != jpi (avoid to extend several times) 163 216 jplus (optional, default=25) : number of points added on the east side of the field 164 217 165 218 ''' 166 jpi = tab.shape[-1] 167 JPLUS = jplus 168 169 try : 170 if type (tab) == xr.core.dataarray.DataArray : math = xr 171 else : math = np 172 except : math = np 173 174 if Lon : xplus = -360.0 175 else : xplus = 0.0 176 177 if math == xr : 178 extend = xr.concat ( (tab[..., 2:jpi] + xplus, tab[..., 2:jplus]), dim=tab.dims[-1] ) 179 if math == np : 180 extend = np.concatenate ( (tab[..., 2:jpi] + xplus, tab[..., 2:jplus]), axis=-1 ) 219 if tab.shape[-1] == 1 : 220 extend = tab 221 222 else : 223 if jpi == None : jpi = tab.shape[-1] 224 225 try : 226 if type (tab) == xr.core.dataarray.DataArray : math = xr 227 else : math = np 228 except : math = np 229 230 if Lon : xplus = -360.0 231 else : xplus = 0.0 232 233 if tab.shape[-1] > jpi : 234 extend = tab 235 else : 236 if nperio == 1 : 237 istart = 0 ; le=jpi+1 ; la=0 238 if nperio > 1 : # OPA case with to halo points for periodicity 239 istart = 1 ; le=jpi-2 ; la=1 # Perfect, except at the pole that should be masked by lbc_plot 240 241 if math == xr : 242 extend = xr.concat ( (tab[..., istart:istart+le+1] + xplus, tab[..., istart+la:jplus]), dim=tab.dims[-1] ) 243 if math == np : 244 extend = np.concatenate ( (tab[..., istart:istart+le+1] + xplus, tab[..., istart+la:jplus]), axis=-1 ) 181 245 182 246 return extend 183 247 248 def orca2reg (ff, lat_name='nav_lat', lon_name='nav_lon', y_name='y', x_name='x') : 249 ''' 250 Assign ORCA dataset on a regular grid. 251 For use in the tropical region. 252 253 Inputs : 254 ff : xarray dataset 255 lat_name, lon_name : name of latitude and longitude 2D field in ff 256 y_name, x_name : namex of dimensions in ff 257 258 Returns : xarray dataset with rectangular grid. Incorrect above 20°N 259 ''' 260 # Compute 1D longitude and latitude 261 (lat, lon) = latlon1D (ff[lat_name], ff[lon_name]) 262 # Assign lon and lat as dimensions of the dataset 263 if y_name in ff.dims : 264 lat = xr.DataArray (lat, coords=[lat,], dims=['lat',]) 265 ff = ff.rename_dims ({y_name: "lat",}).assign_coords (lat=lat) 266 if x_name in ff.dims : 267 lon = xr.DataArray (lon, coords=[lon,], dims=['lon',]) 268 ff = ff.rename_dims ({x_name: "lon",}).assign_coords (lon=lon) 269 # Force dimensions to be in the right order 270 coord_order = ['lat', 'lon'] 271 for dim in [ 'depthw', 'depthv', 'depthu', 'deptht', 'depth', 'z', 272 'time_counter', 'time', 'tbnds', 273 'bnds', 'axis_nbounds', 'two2', 'two1', 'two', 'four',] : 274 if dim in ff.dims : coord_order.insert (0, dim ) 275 276 ff = ff.transpose (*coord_order) 277 return ff 278 279 def lbc_init (ptab, nperio=None, cd_type='T') : 280 """ 281 Prepare for all lbc calls 282 283 Set periodicity on input field 284 nperio : Type of periodicity 285 0 : No periodicity 286 1, 4, 6 : Cyclic on i dimension (generaly longitudes) with 2 points halo 287 2 : Obsolete (was symmetric condition at southern boundary ?) 288 3, 4 : North fold T-point pivot (legacy ORCA2) 289 5, 6 : North fold F-point pivot (ORCA1, ORCA025, ORCA2 with new grid for paleo) 290 cd_type : Grid specification : T, U, V or F 291 292 See NEMO documentation for further details 293 """ 294 jpj, jpi = ptab.shape[-2:] 295 if nperio == None : nperio = __guessNperio__ (jpj, jpi, nperio) 296 297 if nperio not in nperio_valid_range : 298 print ( 'nperio=', nperio, ' is not in the valid range', nperio_valid_range ) 299 sys.exit() 300 301 if cd_type == None and isinstance (ptab, xr.core.dataarray.DataArray) : 302 cd_type = __guessPoint__ (ptab) 303 304 if cd_type not in ['T', 'U', 'V', 'F', 'W', 'F' ]: 305 print ( 'cd_type=', cd_type, ' is not in the list T, U, V, F, W, F' ) 306 sys.exit () 307 308 return jpj, jpi, nperio, cd_type 309 184 310 185 311 def lbc (ptab, nperio=None, cd_type='T', psgn=1.0, nemo_4U_bug=False) : … … 189 315 rank 2 at least : ptab[...., lat, lon] 190 316 nperio : Type of periodicity 191 1, 4, 6 : Cyclic on i dimension (generaly longitudes)192 2 : Obsolete (was symmetric condition at southern boundary ?)193 3, 4 : North fold T-point pivot (legacy ORCA2)194 5, 6 : North fold F-point pivot (ORCA1, ORCA025, ORCA2 with new grid for paleo)195 317 cd_type : Grid specification : T, U, V or F 196 318 psgn : For change of sign for vector components (1 for scalars, -1 for vector components) … … 198 320 See NEMO documentation for further details 199 321 """ 200 201 jpi = ptab.shape[-1] 202 nperio = __guessNperio__ ( jpi, nperio) 203 psgn = ptab.dtype.type(psgn) 204 322 jpj, jpi, nperio, cd_type = lbc_init (ptab, nperio, cd_type) 323 psgn = ptab.dtype.type (psgn) 324 205 325 # 206 326 #> East-West boundary conditions … … 210 330 ptab [..., :, 0] = ptab [..., :, -2] 211 331 ptab [..., :, -1] = ptab [..., :, 1] 212 213 332 # 214 333 #> North-South boundary conditions … … 216 335 if nperio in [3, 4] : # North fold T-point pivot 217 336 if cd_type in [ 'T', 'W' ] : # T-, W-point 218 ptab [..., -1, 1: ] = psgn * ptab [..., -3, -1:0:-1 ]337 ptab [..., -1, 1: ] = psgn * ptab [..., -3, -1:0:-1] 219 338 ptab [..., -1, 0 ] = psgn * ptab [..., -3, 2 ] 220 339 ptab [..., -2, jpi//2: ] = psgn * ptab [..., -2, jpi//2:0:-1 ] 221 340 222 341 if cd_type == 'U' : 223 342 ptab [..., -1, 0:-1 ] = psgn * ptab [..., -3, -1:0:-1 ] 224 343 ptab [..., -1, 0 ] = psgn * ptab [..., -3, 1 ] 225 344 ptab [..., -1, -1 ] = psgn * ptab [..., -3, -2 ] 226 345 227 346 if nemo_4U_bug : 228 347 ptab [..., -2, jpi//2+1:-1] = psgn * ptab [..., -2, jpi//2-2:0:-1] … … 230 349 else : 231 350 ptab [..., -2, jpi//2-1:-1] = psgn * ptab [..., -2, jpi//2:0:-1] 232 351 233 352 if cd_type == 'V' : 234 353 ptab [..., -2, 1: ] = psgn * ptab [..., -3, jpi-1:0:-1 ] 235 354 ptab [..., -1, 1: ] = psgn * ptab [..., -4, -1:0:-1 ] 236 355 ptab [..., -1, 0 ] = psgn * ptab [..., -4, 2 ] 237 356 238 357 if cd_type == 'F' : 239 358 ptab [..., -2, 0:-1 ] = psgn * ptab [..., -3, -1:0:-1 ] … … 241 360 ptab [..., -1, 0 ] = psgn * ptab [..., -4, 1 ] 242 361 ptab [..., -1, -1 ] = psgn * ptab [..., -4, -2 ] 243 362 244 363 if nperio in [5, 6] : # North fold F-point pivot 245 364 if cd_type in ['T', 'W'] : … … 248 367 if cd_type == 'U' : 249 368 ptab [..., -1, 0:-1 ] = psgn * ptab [..., -2, -2::-1 ] 250 ptab [..., -1, -1 ] = psgn * ptab [..., -2, 0 ] 251 369 ptab [..., -1, -1 ] = psgn * ptab [..., -2, 0 ] # Bug ? 370 252 371 if cd_type == 'V' : 253 372 ptab [..., -1, 0: ] = psgn * ptab [..., -3, -1::-1 ] 254 373 ptab [..., -2, jpi//2: ] = psgn * ptab [..., -2, jpi//2-1::-1 ] 255 374 256 375 if cd_type == 'F' : 257 376 ptab [..., -1, 0:-1 ] = psgn * ptab [..., -3, -2::-1 ] 258 377 ptab [..., -1, -1 ] = psgn * ptab [..., -3, 0 ] 259 378 ptab [..., -2, jpi//2:-1] = psgn * ptab [..., -2, jpi//2-2::-1 ] 260 379 261 380 return ptab 262 381 263 def lbc_mask (ptab, nperio=None, cd_type='T', sval= None) :382 def lbc_mask (ptab, nperio=None, cd_type='T', sval=np.nan) : 264 383 # 265 384 """ … … 268 387 rank 2 at least : ptab[...., lat, lon] 269 388 nperio : Type of periodicity 270 1, 4, 6 : Cyclic on i dimension (generaly longitudes)271 2 : Obsolete (was symmetric condition at southern boundary ?)272 3, 4 : North fold T-point pivot (legacy ORCA2)273 5, 6 : North fold F-point pivot (ORCA1, ORCA025, ORCA2 with new grid for paleo)274 389 cd_type : Grid specification : T, U, V or F 275 390 … … 277 392 """ 278 393 279 jpi = ptab.shape[-1] 280 nperio = __guessNperio__ (jpi, nperio) 281 if sval == None : sval = ptab.dtype.type (0) 282 394 jpj, jpi, nperio, cd_type = lbc_init (ptab, nperio, cd_type) 395 283 396 # 284 397 #> East-West boundary conditions … … 303 416 ptab [..., -1, 0 ] = sval 304 417 ptab [..., -2, jpi//2: ] = sval 305 418 306 419 if cd_type == 'U' : 307 ptab [..., -1, 0:-1 ] = sval 308 ptab [..., -1, 0 ] = sval 309 ptab [..., -1, -1 ] = sval 310 ptab [..., -2, jpi//2 :] = sval 311 420 ptab [..., -1, : ] = sval 421 ptab [..., -2, jpi//2: ] = sval 422 312 423 if cd_type == 'V' : 313 ptab [..., -2, 1: ] = sval 314 ptab [..., -1, 1: ] = sval 315 ptab [..., -1, 0 ] = sval 316 424 ptab [..., -2, : ] = sval 425 ptab [..., -1, : ] = sval 426 317 427 if cd_type == 'F' : 318 ptab [..., -2, 0:-1 ] = sval 319 ptab [..., -1, 0:-1 ] = sval 320 ptab [..., -1, 0 ] = sval 321 ptab [..., -1, -1 ] = sval 322 428 ptab [..., -2, : ] = sval 429 ptab [..., -1, : ] = sval 430 323 431 if nperio in [5, 6] : # North fold F-point pivot 324 432 if cd_type in ['T', 'W'] : … … 340 448 return ptab 341 449 450 451 def lbc_plot (ptab, nperio=None, cd_type='T', psgn=1.0, sval=np.nan) : 452 """ 453 Set periodicity on input field, adapted for plotting for any cartopy projection 454 ptab : Input array (works 455 rank 2 at least : ptab[...., lat, lon] 456 nperio : Type of periodicity 457 cd_type : Grid specification : T, U, V or F 458 psgn : For change of sign for vector components (1 for scalars, -1 for vector components) 459 460 See NEMO documentation for further details 461 """ 462 463 jpj, jpi, nperio, cd_type = lbc_init (ptab, nperio, cd_type) 464 psgn = ptab.dtype.type (psgn) 465 466 # 467 #> East-West boundary conditions 468 # ------------------------------ 469 if nperio in [1, 4, 6] : 470 # ... cyclic 471 ptab [..., :, 0] = ptab [..., :, -2] 472 ptab [..., :, -1] = ptab [..., :, 1] 473 474 # 475 #> North-South boundary conditions 476 # -------------------------------- 477 if nperio in [3, 4] : # North fold T-point pivot 478 if cd_type in [ 'T', 'W' ] : # T-, W-point 479 ptab [..., -1, :] = sval 480 ptab [..., -2, :jpi//2] = sval 481 #ptab [..., -2, :] = sval 482 483 if cd_type == 'U' : 484 ptab [..., -1, : ] = sval 485 486 if cd_type == 'V' : 487 ptab [..., -2, : ] = sval 488 ptab [..., -1, : ] = sval 489 490 if cd_type == 'F' : 491 ptab [..., -2, : ] = sval 492 ptab [..., -1, : ] = sval 493 494 if nperio in [5, 6] : # North fold F-point pivot 495 if cd_type in ['T', 'W'] : 496 ptab [..., -1, : ] = sval 497 498 if cd_type == 'U' : 499 ptab [..., -1, : ] = sval 500 501 if cd_type == 'V' : 502 ptab [..., -1, : ] = sval 503 ptab [..., -2, jpi//2: ] = sval 504 505 if cd_type == 'F' : 506 ptab [..., -1, : ] = sval 507 ptab [..., -2, jpi//2:-1] = sval 508 509 return ptab 510 511 def geo2oce (pxx, pyy, pzz, glam, gphi) : 512 ''' 513 Change vector from geocentric to east/north 514 515 Input 516 pxx, pyy, pzz : components on the geoce,tric system 517 glam, gphi : longitue and latitude of the points 518 519 ''' 520 gsinlon = np.sin (rad * glam) 521 gcoslon = np.cos (rad * glam) 522 gsinlat = np.sin (rad * gphi) 523 gcoslat = np.cos (rad * gphi) 524 525 pte = - pxx * gsinlon + pyy * gcoslon 526 ptn = - pxx * gcoslon * gsinlat - pyy * gsinlon * gsinlat + pzz * gcoslat 527 528 return pte, ptn 529 530 def oce2geo (pte, ptn, glam, gphi) : 531 ##---------------------------------------------------------------------- 532 ## *** ROUTINE oce2geo *** 533 ## 534 ## ** Purpose : 535 ## 536 ## ** Method : Change vector from east/north to geocentric 537 ## 538 ## History : 539 ## # (A. Caubel) oce2geo - Original code 540 ##---------------------------------------------------------------------- 541 gsinlon = np.sin (rad * glam) 542 gcoslon = np.cos (rad * glam) 543 gsinlat = np.sin (rad * gphi) 544 gcoslat = np.cos (rad * gphi) 545 546 pxx = - pte * gsinlon - ptn * gcoslon * gsinlat 547 pyy = pte * gcoslon - ptn * gsinlon * gsinlat 548 pzz = ptn * gcoslat 549 550 return pxx, pyy, pzz 551 342 552 ## =========================================================================== 343 553 ##
Note: See TracChangeset
for help on using the changeset viewer.