Changeset 6245 for TOOLS/MOSAIX
- Timestamp:
- 09/20/22 11:40:57 (20 months ago)
- File:
-
- 1 edited
Legend:
- Unmodified
- Added
- Removed
-
TOOLS/MOSAIX/nemo.py
r6190 r6245 2 2 ## =========================================================================== 3 3 ## 4 ## MOSAIX is under CeCILL_V2 licence. See "Licence_CeCILL_V2-en.txt" 5 ## file for an english version of the licence and 6 ## "Licence_CeCILL_V2-fr.txt" for a french version. 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". 7 9 ## 8 ## Permission is hereby granted, free of charge, to any person or 9 ## organization obtaining a copy of the software and accompanying 10 ## documentation covered by this license (the "Software") to use, 11 ## reproduce, display, distribute, execute, and transmit the 12 ## Software, and to prepare derivative works of the Software, and to 13 ## permit third-parties to whom the Software is furnished to do so, 14 ## all subject to the following: 15 ## 16 ## Warning, to install, configure, run, use any of MOSAIX software or 17 ## to read the associated documentation you'll need at least one (1) 18 ## brain in a reasonably working order. Lack of this implement will 19 ## void any warranties (either express or implied). Authors assumes 20 ## no responsability for errors, omissions, data loss, or any other 21 ## consequences caused directly or indirectly by the usage of his 22 ## software by incorrectly or partially configured 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. 23 18 ## 24 19 ## =========================================================================== … … 37 32 __HeadURL = "$HeadURL$" 38 33 39 import sys,numpy as np34 import numpy as np 40 35 try : import xarray as xr 36 except ImportError : pass 37 38 try : import f90nml 39 except : pass 40 41 try : from sklearn.impute import SimpleImputer 42 except : pass 43 44 try : import numba 41 45 except : pass 42 46 … … 62 66 tList = [ 't', 'T', 'time' , ] 63 67 64 ## SVN information 65 __Author__ = "$Author$" 66 __Date__ = "$Date$" 67 __Revision__ = "$Revision$" 68 __Id__ = "$Id$" 69 __HeadURL = "$HeadURL$" 70 71 def __math__ (tab, default=None) : 72 math = default 68 ## =========================================================================== 69 def __mmath__ (tab, default=None) : 70 mmath = default 73 71 try : 74 if type (tab) == xr.core.dataarray.DataArray : m ath = xr72 if type (tab) == xr.core.dataarray.DataArray : mmath = xr 75 73 except : 76 74 pass 77 75 78 76 try : 79 if type (tab) == np.ndarray : m ath = np77 if type (tab) == np.ndarray : mmath = np 80 78 except : 81 79 pass 82 80 83 return math 84 85 def __guessNperio__ (jpj, jpi, nperio=None) : 81 return mmath 82 83 84 def __guessNperio__ (jpj, jpi, nperio=None, out='nperio') : 86 85 ''' 87 86 Tries to guess the value of nperio (periodicity parameter. See NEMO documentation for details) … … 93 92 ''' 94 93 if nperio == None : 94 nperio = __guessConfig__ (jpj, jpi, nperio=None, out='nperio') 95 96 return nperio 97 98 def __guessConfig__ (jpj, jpi, nperio=None, config=None, out='nperio') : 99 ''' 100 Tries to guess the value of nperio (periodicity parameter. See NEMO documentation for details) 101 102 Inputs 103 jpj : number of latitudes 104 jpi : number of longitudes 105 nperio : periodicity parameter 106 ''' 107 if nperio == None : 95 108 ## Values for NEMO version < 4.2 96 if jpi == 182 : 109 if jpj == 149 and jpi == 182 : 110 config = 'ORCA2.3' 97 111 nperio = 4 # ORCA2. We choose legacy orca2. 98 112 Iperio = 1 ; Jperio = 0 ; NFold = 1 ; NFtype = 'T' 99 if jpi == 362 : # eORCA1. 113 if jpj == 332 and jpi == 362 : # eORCA1. 114 config = 'eORCA1.2' 100 115 nperio = 6 101 116 Iperio = 1 ; Jperio = 0 ; NFold = 1 ; NFtype = 'F' 102 117 if jpi == 1442 : # ORCA025. 118 config = 'ORCA025' 103 119 nperio = 6 104 120 Iperio = 1 ; Jperio = 0 ; NFold = 1 ; NFtype = 'F' 105 if jpj == 149 : # ORCA2. We choose legacy orca2.106 nperio = 4107 Iperio = 1 ; Jperio = 0 ; NFold = 1 ; NFtype = 'F'108 121 if jpj == 294 : # ORCA1 122 config = 'ORCA1' 109 123 nperio = 6 110 Iperio = 1 ; Jperio = 0 ; NFold = 1 ; NFtype = 'F'111 if jpj == 332 : # eORCA1.112 nperio = 6113 124 Iperio = 1 ; Jperio = 0 ; NFold = 1 ; NFtype = 'F' 114 125 115 126 ## Values for NEMO version >= 4.2. No more halo points 116 if jpi == 180 : 127 if jpj == 148 and jpi == 180 : 128 config = 'ORCA2.4' 117 129 nperio = 4.2 # ORCA2. We choose legacy orca2. 118 130 Iperio = 1 ; Jperio = 0 ; NFold = 1 ; NFtype = 'F' 119 if jpi == 360 : # eORCA1. 131 if jpj == 331 and jpi == 360 : # eORCA1. 132 config = 'eORCA1.4' 120 133 nperio = 6.2 121 134 Iperio = 1 ; Jperio = 0 ; NFold = 1 ; NFtype = 'F' 122 135 if jpi == 1440 : # ORCA025. 136 config = 'ORCA025' 123 137 nperio = 6.2 124 138 Iperio = 1 ; Jperio = 0 ; NFold = 1 ; NFtype = 'F' 125 126 if jpj == 148 : nperio = 4.2 # ORCA2. We choose legacy orca2. 127 if jpj == 330 : nperio = 6.2 # eORCA1. 128 139 129 140 if nperio == None : 130 sys.exit('in nemo module : nperio not found, and cannot by guessed')141 raise Exception ('in nemo module : nperio not found, and cannot by guessed') 131 142 else : 132 143 if nperio in nperio_valid_range : 133 print ('nperio set as {:} (deduced from jp i={:d})'.format (nperio, jpi))144 print ('nperio set as {:} (deduced from jpj={:d} jpi={:d})'.format (nperio, jpj, jpi)) 134 145 else : 135 sys.exit ('nperio set as {:} (deduced from jpi={:d}) : nemo.py is not ready for this value'.format (nperio, jpi)) 136 137 return nperio 138 146 raise ValueError ('nperio set as {:} (deduced from jpi={:d}) : nemo.py is not ready for this value'.format (nperio, jpi)) 147 148 if out == 'nperio' : return nperio 149 if out == 'config' : return config 150 if out == 'perio' : return Iperio, Jperio, NFold, NFtype 151 if out in ['full', 'all'] : return {'nperio':nperio, 'Iperio':Iperio, 'Jperio':Jperio, 'NFold':NFold, 'NFtype':NFtype} 152 139 153 def __guessPoint__ (ptab) : 140 154 ''' … … 149 163 ''' 150 164 gP = None 151 m ath = __math__ (ptab)152 if m ath == xr :165 mmath = __mmath__ (ptab) 166 if mmath == xr : 153 167 if 'x_c' in ptab.dims and 'y_c' in ptab.dims : gP = 'T' 154 168 if 'x_f' in ptab.dims and 'y_c' in ptab.dims : gP = 'U' … … 162 176 163 177 if gP == None : 164 sys.exit('in nemo module : cd_type not found, and cannot by guessed')178 raise Exception ('in nemo module : cd_type not found, and cannot by guessed') 165 179 else : 166 180 print ('Grid set as', gP, 'deduced from dims ', ptab.dims) 167 181 return gP 168 182 else : 169 sys.exit ('in nemo module : cd_type not found, input is not an xarray data') 170 #return gP 183 raise Exception ('in nemo module : cd_type not found, input is not an xarray data') 184 185 def lbc_diag (nperio) : 186 lperio = nperio ; aperio = False 187 if nperio == 4.2 : 188 lperio = 4 ; aperio = True 189 if nperio == 6.2 : 190 lperio = 6 ; aperio = True 191 192 return lperio, aperio 171 193 172 194 def __findAxis__ (tab, axis='z') : … … 174 196 Find number and name of the requested axis 175 197 ''' 176 m ath = __math__ (tab)198 mmath = __mmath__ (tab) 177 199 ix = None ; ax = None 178 200 … … 198 220 unList = [ 'second', 'minute', 'hour', 'day', 'month' ] 199 221 200 if m ath == xr :222 if mmath == xr : 201 223 for Name in axList : 202 224 try : … … 222 244 223 245 return ix, ax 224 246 247 #@numba.jit(forceobj=True) 225 248 def fixed_lon (lon, center_lon=0.0) : 226 249 ''' … … 232 255 Designed by Phil Pelson. See https://gist.github.com/pelson/79cf31ef324774c97ae7 233 256 ''' 234 m ath = __math__ (lon)257 mmath = __mmath__ (lon) 235 258 236 259 fixed_lon = lon.copy () 237 260 238 fixed_lon = m ath.where (fixed_lon > center_lon+180., fixed_lon-360.0, fixed_lon)239 fixed_lon = m ath.where (fixed_lon < center_lon-180., fixed_lon+360.0, fixed_lon)261 fixed_lon = mmath.where (fixed_lon > center_lon+180., fixed_lon-360.0, fixed_lon) 262 fixed_lon = mmath.where (fixed_lon < center_lon-180., fixed_lon+360.0, fixed_lon) 240 263 241 264 for i, start in enumerate (np.argmax (np.abs (np.diff (fixed_lon, axis=-1)) > 180., axis=-1)) : 242 fixed_lon [..., i, start+1:] += 360 265 fixed_lon [..., i, start+1:] += 360. 243 266 244 267 # Special case for eORCA025 … … 246 269 if fixed_lon.shape [-1] == 1440 : fixed_lon [..., -1, :] = fixed_lon [..., -2, :] 247 270 248 if fixed_lon.min () > center_lon : fixed_lon = fixed_lon-360.0 271 if fixed_lon.min () > center_lon : fixed_lon += -360.0 272 if fixed_lon.max () < center_lon : fixed_lon += 360.0 273 274 if fixed_lon.min () < center_lon-360.0 : fixed_lon += 360.0 275 if fixed_lon.max () > center_lon+360.0 : fixed_lon += -360.0 249 276 250 277 return fixed_lon 251 278 279 #@numba.jit(forceobj=True) 280 def fill_empty (ztab, sval=np.nan, transpose=False) : 281 ''' 282 Fill values 283 284 Useful when NEMO has run with no wet points options : 285 some parts of the domain, with no ocean points, has no 286 lon/lat values 287 ''' 288 mmath = __mmath__ (ztab) 289 290 imp = SimpleImputer (missing_values=sval, strategy='mean') 291 if transpose : 292 imp.fit (ztab.T) 293 ptab = imp.transform (ztab.T).T 294 else : 295 imp.fit (ztab) 296 ptab = imp.transform (ztab) 297 298 if mmath == xr : 299 ptab = xr.DataArray (ptab, dims=ztab.dims, coords=ztab.coords) 300 ptab.attrs = ztab.attrs 301 302 return ptab 303 304 #@numba.jit(forceobj=True) 305 def fill_lonlat (lon, lat, sval=-1) : 306 ''' 307 Fill longitude/latitude values 308 309 Useful when NEMO has run with no wet points options : 310 some parts of the domain, with no ocean points, as no 311 lon/lat values 312 ''' 313 mmath = __mmath__ (lon) 314 315 imp = SimpleImputer (missing_values=sval, strategy='mean') 316 imp.fit (lon) 317 plon = imp.transform (lon) 318 imp.fit (lat.T) 319 plat = imp.transform (lat.T).T 320 321 if mmath == xr : 322 plon = xr.DataArray (plon, dims=lon.dims, coords=lon.coords) 323 plat = xr.DataArray (plat, dims=lat.dims, coords=lat.coords) 324 plon.attrs = lon.attrs ; plat.attrs = lat.attrs 325 326 plon = fixed_lon (plon) 327 328 return plon, plat 329 330 #@numba.jit(forceobj=True) 252 331 def jeq (lat) : 253 332 ''' … … 256 335 lat : latitudes of the grid. At least 2D. 257 336 ''' 258 m ath = __math__ (lat)337 mmath = __mmath__ (lat) 259 338 ix, ax = __findAxis__ (lat, 'x') 260 339 iy, ay = __findAxis__ (lat, 'y') 261 340 262 if m ath == xr :341 if mmath == xr : 263 342 jeq = int ( np.mean ( np.argmin (np.abs (np.float64 (lat)), axis=iy) ) ) 264 343 else : … … 266 345 return jeq 267 346 347 #@numba.jit(forceobj=True) 268 348 def lon1D (lon, lat=None) : 269 349 ''' … … 273 353 lat (optionnal) : latitudes of the grid 274 354 ''' 355 mmath = __mmath__ (lon) 275 356 if np.max (lat) != None : 276 357 je = jeq (lat) 277 lon1D = lon [..., je, :]358 lon1D = lon.copy() [..., je, :] 278 359 else : 279 360 jpj, jpi = lon.shape [-2:] 280 lon1D = lon [..., jpj//3, :]361 lon1D = lon.copy() [..., jpj//3, :] 281 362 282 363 start = np.argmax (np.abs (np.diff (lon1D, axis=-1)) > 180.0, axis=-1) 283 364 lon1D [..., start+1:] += 360 365 366 if mmath == xr : 367 lon1D.attrs = lon.attrs 368 lon1D = lon1D.assign_coords ( {'x':lon1D} ) 284 369 285 370 return lon1D 286 371 372 #@numba.jit(forceobj=True) 287 373 def latreg (lat, diff=0.1) : 288 374 ''' 289 Returns maximum j index where gridlines are along latitude in the northern hemisphere375 Returns maximum j index where gridlines are along latitudes in the northern hemisphere 290 376 291 377 lat : latitudes of the grid (2D) 292 378 diff [optional] : tolerance 293 379 ''' 294 m ath = __math__ (lat)380 mmath = __mmath__ (lat) 295 381 if diff == None : 296 382 dy = np.float64 (np.mean (np.abs (lat - np.roll (lat,shift=1,axis=-2, roll_coords=False)))) … … 304 390 return jreg, latreg 305 391 392 #@numba.jit(forceobj=True) 306 393 def lat1D (lat) : 307 394 ''' … … 310 397 lat : latitudes of the grid (2D) 311 398 ''' 312 m ath = __math__ (lat)399 mmath = __mmath__ (lat) 313 400 jpj, jpi = lat.shape[-2:] 314 401 … … 326 413 yrange = 90. -lat_reg 327 414 328 lat1D = math.where (lat_ave<lat_reg, lat_ave, lat_reg + yrange * (np.arange(jpj)-jreg)/(jpj-jreg-1)) 329 330 if math == xr : lat1D.attrs = lat.attrs 415 lat1D = mmath.where (lat_ave<lat_reg, lat_ave, lat_reg + yrange * (np.arange(jpj)-jreg)/(jpj-jreg-1)) 416 417 if mmath == xr : 418 lat1D.attrs = lat.attrs 419 lat1D = lat1D.assign_coords ( {'y':lat1D} ) 331 420 332 421 return lat1D 333 422 423 #@numba.jit(forceobj=True) 334 424 def latlon1D (lat, lon) : 335 425 ''' … … 340 430 return lat1D (lat), lon1D (lon, lat) 341 431 432 ##@numba.jit(forceobj=True) 342 433 def mask_lonlat (ptab, x0, x1, y0, y1, lon, lat, sval=np.nan) : 343 m ath = __math__ (ptab)434 mmath = __mmath__ (ptab) 344 435 try : 345 436 lon = lon.copy().to_masked_array() … … 350 441 np.logical_or (np.logical_or (np.logical_and(lon>x0, lon<x1), np.logical_and(lon+360>x0, lon+360<x1)), 351 442 np.logical_and(lon-360>x0, lon-360<x1))) 352 tab = m ath.where (mask, ptab, np.nan)443 tab = mmath.where (mask, ptab, np.nan) 353 444 354 445 return tab 355 446 447 #@numba.jit(forceobj=True) 356 448 def extend (tab, Lon=False, jplus=25, jpi=None, nperio=4) : 357 449 ''' … … 366 458 367 459 ''' 368 m ath = __math__ (tab)460 mmath = __mmath__ (tab) 369 461 370 462 if tab.shape[-1] == 1 : extend = tab … … 386 478 istart = 1 ; le=jpi-2 ; la=1 # Perfect, except at the pole that should be masked by lbc_plot 387 479 388 if m ath == xr :480 if mmath == xr : 389 481 extend = np.concatenate ((tab.values[..., istart :istart+le+1 ] + xplus, 390 482 tab.values[..., istart+la:istart+la+jplus] ), axis=-1) … … 402 494 def orca2reg (ff, lat_name='nav_lat', lon_name='nav_lon', y_name='y', x_name='x') : 403 495 ''' 404 Assign ORCA dataset on a regular grid.496 Assign an ORCA dataset on a regular grid. 405 497 For use in the tropical region. 406 498 407 499 Inputs : 408 500 ff : xarray dataset 409 501 lat_name, lon_name : name of latitude and longitude 2D field in ff 410 502 y_name, x_name : namex of dimensions in ff 411 412 Returns : xarray dataset with rectangular grid. Incorrect above 20°N503 504 Returns : xarray dataset with rectangular grid. Incorrect above 20°N 413 505 ''' 414 506 # Compute 1D longitude and latitude 415 507 (lat, lon) = latlon1D (ff[lat_name], ff[lon_name]) 508 416 509 # Assign lon and lat as dimensions of the dataset 417 510 if y_name in ff.dims : … … 450 543 451 544 if nperio not in nperio_valid_range : 452 print ('nperio=', nperio, ' is not in the valid range', nperio_valid_range) 453 sys.exit () 545 raise Exception ('nperio=', nperio, ' is not in the valid range', nperio_valid_range) 454 546 455 547 return jpj, jpi, nperio 456 548 457 549 #@numba.jit(forceobj=True) 458 550 def lbc (ptab, nperio=None, cd_type='T', psgn=1.0, nemo_4U_bug=False) : 459 551 ''' … … 468 560 jpj, jpi, nperio = lbc_init (ptab, nperio) 469 561 psgn = ptab.dtype.type (psgn) 470 math = __math__ (ptab) 471 472 if math == xr : ztab = ptab.values.copy () 473 else : ztab = ptab.copy () 562 mmath = __mmath__ (ptab) 563 564 if mmath == xr : ztab = ptab.values.copy () 565 else : ztab = ptab.copy () 566 474 567 # 475 568 #> East-West boundary conditions … … 548 641 ztab [..., :, -1] = ztab [..., :, 1] 549 642 550 if m ath == xr :643 if mmath == xr : 551 644 ztab = xr.DataArray ( ztab, dims=ptab.dims, coords=ptab.coords) 645 ztab.attrs = ptab.attrs 552 646 553 647 return ztab 554 648 649 #@numba.jit(forceobj=True) 555 650 def lbc_mask (ptab, nperio=None, cd_type='T', sval=np.nan) : 556 651 # … … 632 727 return ztab 633 728 729 #@numba.jit(forceobj=True) 634 730 def lbc_plot (ptab, nperio=None, cd_type='T', psgn=1.0, sval=np.nan) : 635 731 ''' … … 707 803 return ztab 708 804 709 def lbc_add (ptab, nperio=None, cd_type =None) : 805 #@numba.jit(forceobj=True) 806 def lbc_add (ptab, nperio=None, cd_type=None, psgn=1, sval=None) : 710 807 ''' 711 808 Handle NEMO domain changes between NEMO 4.0 to NEMO 4.2 712 Peridodicity halo has been remove s809 Peridodicity halo has been removed 713 810 This routine adds the halos if needed 714 811 … … 719 816 See NEMO documentation for further details 720 817 ''' 721 818 mmath = __mmath__ (ptab) 722 819 jpj, jpi, nperio = lbc_init (ptab, nperio) 723 820 … … 725 822 726 823 if nperio == 4.2 or nperio == 6.2 : 727 math = __math__ (ptab)728 824 729 825 ext_shape = t_shape … … 731 827 ext_shape[-2] = ext_shape[-2] + 1 732 828 733 if math == xr : ptab_ext = xr.DataArray (np.empty (ext_shape), dims=ptab.dims) 734 else : ptab_ext = np.empty (ext_shape) 829 if mmath == xr : 830 ptab_ext = xr.DataArray (np.zeros (ext_shape), dims=ptab.dims) 831 ptab_ext.values[..., :-1, 1:-1] = ptab.values.copy () 832 else : 833 ptab_ext = np.zeros (ext_shape) 834 ptab_ext[..., :-1, 1:-1] = ptab.copy () 735 835 736 ptab_ext[..., :-1, 1:-1] = ptab 737 738 if nperio == 4.2 : ptab_ext = lbc (ptab_ext, 4, cd_type) 739 if nperio == 6.2 : ptab_ext = lbc (ptab_ext, 6, cd_type) 740 741 if math == xr : 742 for attr in ptab.attrs : 743 ptab_ext.attrs [attr] = ptab.attrs [attr] 744 745 else : ptab_ext = ptab 746 747 return ptab 748 749 def lbc_del (ptab, nperio=None) : 836 #if sval != None : ptab_ext[..., 0, :] = sval 837 838 if nperio == 4.2 : ptab_ext = lbc (ptab_ext, nperio=4, cd_type=cd_type, psgn=psgn) 839 if nperio == 6.2 : ptab_ext = lbc (ptab_ext, nperio=6, cd_type=cd_type, psgn=psgn) 840 841 if mmath == xr : 842 ptab_ext.attrs = ptab.attrs 843 844 else : ptab_ext = lbc (ptab, nperio=nperio, cd_type=cd_type, psgn=psgn) 845 846 return ptab_ext 847 848 def lbc_del (ptab, nperio=None, cd_type='T', psgn=1) : 750 849 ''' 751 850 Handle NEMO domain changes between NEMO 4.0 to NEMO 4.2 752 Peri dodicity halo has been removes851 Periodicity halo has been removed 753 852 This routine removes the halos if needed 754 853 … … 762 861 jpj, jpi, nperio = lbc_init (ptab, nperio) 763 862 764 if nperio == 4 or nperio == 6:765 return ptab[..., :-1, 1:-1]863 if nperio == 4.2 or nperio == 6.2 : 864 return lbc (ptab[..., :-1, 1:-1], nperio=nperio, cd_type=cd_type, psgn=psgn) 766 865 else : 767 866 return ptab 768 867 868 #@numba.jit(forceobj=True) 769 869 def lbc_index (jj, ii, jpj, jpi, nperio=None, cd_type='T') : 770 870 ''' … … 784 884 jy = jj + 1 ; ix = ii + 1 785 885 786 m ath = __math__ (jj)787 if m ath == None :math=np886 mmath = __mmath__ (jj) 887 if mmath == None : mmath=np 788 888 789 889 # … … 792 892 if nperio in [1, 4, 6] : 793 893 #... cyclic 794 ix = m ath.where (ix==jpi, 2 , ix)795 ix = m ath.where (ix== 1 ,jpi-1, ix)894 ix = mmath.where (ix==jpi, 2 , ix) 895 ix = mmath.where (ix== 1 ,jpi-1, ix) 796 896 797 897 # 798 898 def modIJ (cond, jy_new, ix_new) : 799 jy_r = m ath.where (cond, jy_new, jy)800 ix_r = m ath.where (cond, ix_new, ix)899 jy_r = mmath.where (cond, jy_new, jy) 900 ix_r = mmath.where (cond, ix_new, ix) 801 901 return jy_r, ix_r 802 902 # … … 907 1007 else : jpj = len(lat_grid) ; jpi=len(lon_grid) 908 1008 909 m ath = __math__ (lat_grid)1009 mmath = __mmath__ (lat_grid) 910 1010 911 1011 # Compute distance from the point to all grid points (in radian) … … 943 1043 def clo_lon (lon, lon0) : 944 1044 '''Choose closest to lon0 longitude, adding or substacting 360° if needed''' 945 m ath = __math__ (lon, np)1045 mmath = __mmath__ (lon, np) 946 1046 947 1047 clo_lon = lon 948 clo_lon = m ath.where (clo_lon > lon0 + 180., clo_lon-360., clo_lon)949 clo_lon = m ath.where (clo_lon < lon0 - 180., clo_lon+360., clo_lon)950 clo_lon = m ath.where (clo_lon > lon0 + 180., clo_lon-360., clo_lon)951 clo_lon = m ath.where (clo_lon < lon0 - 180., clo_lon+360., clo_lon)1048 clo_lon = mmath.where (clo_lon > lon0 + 180., clo_lon-360., clo_lon) 1049 clo_lon = mmath.where (clo_lon < lon0 - 180., clo_lon+360., clo_lon) 1050 clo_lon = mmath.where (clo_lon > lon0 + 180., clo_lon-360., clo_lon) 1051 clo_lon = mmath.where (clo_lon < lon0 - 180., clo_lon+360., clo_lon) 952 1052 if clo_lon.shape == () : clo_lon = clo_lon.item () 953 1053 return clo_lon … … 955 1055 def angle_full (glamt, gphit, glamu, gphiu, glamv, gphiv, glamf, gphif, nperio=None) : 956 1056 '''Compute sinus and cosinus of model line direction with respect to east''' 957 math = __math__ (glamt) 1057 mmath = __mmath__ (glamt) 1058 1059 zlamt = lbc_add (glamt, nperio, 'T', 1.) 1060 zphit = lbc_add (gphit, nperio, 'T', 1.) 1061 zlamu = lbc_add (glamu, nperio, 'U', 1.) 1062 zphiu = lbc_add (gphiu, nperio, 'U', 1.) 1063 zlamv = lbc_add (glamv, nperio, 'V', 1.) 1064 zphiv = lbc_add (gphiv, nperio, 'V', 1.) 1065 zlamf = lbc_add (glamf, nperio, 'F', 1.) 1066 zphif = lbc_add (gphif, nperio, 'F', 1.) 958 1067 959 1068 # north pole direction & modulous (at T-point) 960 zxnpt = 0. - 2.0 * np.cos (rad* glamt) * np.tan (rpi/4.0 - rad*gphit/2.0)961 zynpt = 0. - 2.0 * np.sin (rad* glamt) * np.tan (rpi/4.0 - rad*gphit/2.0)1069 zxnpt = 0. - 2.0 * np.cos (rad*zlamt) * np.tan (rpi/4.0 - rad*zphit/2.0) 1070 zynpt = 0. - 2.0 * np.sin (rad*zlamt) * np.tan (rpi/4.0 - rad*zphit/2.0) 962 1071 znnpt = zxnpt*zxnpt + zynpt*zynpt 963 1072 964 1073 # north pole direction & modulous (at U-point) 965 zxnpu = 0. - 2.0 * np.cos (rad* glamu) * np.tan (rpi/4.0 - rad*gphiu/2.0)966 zynpu = 0. - 2.0 * np.sin (rad* glamu) * np.tan (rpi/4.0 - rad*gphiu/2.0)1074 zxnpu = 0. - 2.0 * np.cos (rad*zlamu) * np.tan (rpi/4.0 - rad*zphiu/2.0) 1075 zynpu = 0. - 2.0 * np.sin (rad*zlamu) * np.tan (rpi/4.0 - rad*zphiu/2.0) 967 1076 znnpu = zxnpu*zxnpu + zynpu*zynpu 968 1077 969 1078 # north pole direction & modulous (at V-point) 970 zxnpv = 0. - 2.0 * np.cos (rad* glamv) * np.tan (rpi/4.0 - rad*gphiv/2.0)971 zynpv = 0. - 2.0 * np.sin (rad* glamv) * np.tan (rpi/4.0 - rad*gphiv/2.0)1079 zxnpv = 0. - 2.0 * np.cos (rad*zlamv) * np.tan (rpi/4.0 - rad*zphiv/2.0) 1080 zynpv = 0. - 2.0 * np.sin (rad*zlamv) * np.tan (rpi/4.0 - rad*zphiv/2.0) 972 1081 znnpv = zxnpv*zxnpv + zynpv*zynpv 973 1082 974 1083 # north pole direction & modulous (at F-point) 975 zxnpf = 0. - 2.0 * np.cos( rad* glamf ) * np.tan ( rpi/4. - rad*gphif/2. )976 zynpf = 0. - 2.0 * np.sin( rad* glamf ) * np.tan ( rpi/4. - rad*gphif/2. )1084 zxnpf = 0. - 2.0 * np.cos( rad*zlamf ) * np.tan ( rpi/4. - rad*zphif/2. ) 1085 zynpf = 0. - 2.0 * np.sin( rad*zlamf ) * np.tan ( rpi/4. - rad*zphif/2. ) 977 1086 znnpf = zxnpf*zxnpf + zynpf*zynpf 978 1087 979 1088 # j-direction: v-point segment direction (around T-point) 980 zlam = glamv981 zphi = gphiv982 zlan = np.roll ( glamv, axis=-2, shift=1) # glamv (ji,jj-1)983 zphh = np.roll ( gphiv, axis=-2, shift=1) # gphiv (ji,jj-1)1089 zlam = zlamv 1090 zphi = zphiv 1091 zlan = np.roll ( zlamv, axis=-2, shift=1) # glamv (ji,jj-1) 1092 zphh = np.roll ( zphiv, axis=-2, shift=1) # gphiv (ji,jj-1) 984 1093 zxvvt = 2.0 * np.cos ( rad*zlam ) * np.tan ( rpi/4. - rad*zphi/2. ) \ 985 1094 - 2.0 * np.cos ( rad*zlan ) * np.tan ( rpi/4. - rad*zphh/2. ) … … 989 1098 990 1099 # j-direction: f-point segment direction (around u-point) 991 zlam = glamf992 zphi = gphif993 zlan = np.roll ( glamf, axis=-2, shift=1) # glamf (ji,jj-1)994 zphh = np.roll ( gphif, axis=-2, shift=1) # gphif (ji,jj-1)1100 zlam = zlamf 1101 zphi = zphif 1102 zlan = np.roll (zlamf, axis=-2, shift=1) # glamf (ji,jj-1) 1103 zphh = np.roll (zphif, axis=-2, shift=1) # gphif (ji,jj-1) 995 1104 zxffu = 2.0 * np.cos ( rad*zlam ) * np.tan ( rpi/4. - rad*zphi/2. ) \ 996 1105 - 2.0 * np.cos ( rad*zlan ) * np.tan ( rpi/4. - rad*zphh/2. ) … … 1000 1109 1001 1110 # i-direction: f-point segment direction (around v-point) 1002 zlam = glamf1003 zphi = gphif1004 zlan = np.roll ( glamf, axis=-1, shift=1) # glamf (ji-1,jj)1005 zphh = np.roll ( gphif, axis=-1, shift=1) # gphif (ji-1,jj)1111 zlam = zlamf 1112 zphi = zphif 1113 zlan = np.roll (zlamf, axis=-1, shift=1) # glamf (ji-1,jj) 1114 zphh = np.roll (zphif, axis=-1, shift=1) # gphif (ji-1,jj) 1006 1115 zxffv = 2.0 * np.cos ( rad*zlam ) * np.tan ( rpi/4. - rad*zphi/2. ) \ 1007 1116 - 2.0 * np.cos ( rad*zlan ) * np.tan ( rpi/4. - rad*zphh/2. ) … … 1011 1120 1012 1121 # j-direction: u-point segment direction (around f-point) 1013 zlam = np.roll ( glamu, axis=-2, shift=-1) # glamu (ji,jj+1)1014 zphi = np.roll ( gphiu, axis=-2, shift=-1) # gphiu (ji,jj+1)1015 zlan = glamu1016 zphh = gphiu1122 zlam = np.roll (zlamu, axis=-2, shift=-1) # glamu (ji,jj+1) 1123 zphi = np.roll (zphiu, axis=-2, shift=-1) # gphiu (ji,jj+1) 1124 zlan = zlamu 1125 zphh = zphiu 1017 1126 zxuuf = 2. * np.cos ( rad*zlam ) * np.tan ( rpi/4. - rad*zphi/2. ) \ 1018 1127 - 2. * np.cos ( rad*zlan ) * np.tan ( rpi/4. - rad*zphh/2. ) … … 1035 1144 gcosv =-( zxnpv*zyffv - zynpv*zxffv ) / znffv # (caution, rotation of 90 degres) 1036 1145 1037 gsint = lbc (gsint, cd_type='T', nperio=nperio, psgn=-1.) 1038 gcost = lbc (gcost, cd_type='T', nperio=nperio, psgn=-1.) 1039 gsinu = lbc (gsinu, cd_type='U', nperio=nperio, psgn=-1.) 1040 gcosu = lbc (gcosu, cd_type='U', nperio=nperio, psgn=-1.) 1041 gsinv = lbc (gsinv, cd_type='V', nperio=nperio, psgn=-1.) 1042 gcosv = lbc (gcosv, cd_type='V', nperio=nperio, psgn=-1.) 1043 gsinf = lbc (gsinf, cd_type='F', nperio=nperio, psgn=-1.) 1044 gcosf = lbc (gcosf, cd_type='F', nperio=nperio, psgn=-1.) 1045 1046 if math == xr : 1146 #gsint = lbc (gsint, cd_type='T', nperio=nperio, psgn=-1.) 1147 #gcost = lbc (gcost, cd_type='T', nperio=nperio, psgn=-1.) 1148 #gsinu = lbc (gsinu, cd_type='U', nperio=nperio, psgn=-1.) 1149 #gcosu = lbc (gcosu, cd_type='U', nperio=nperio, psgn=-1.) 1150 #gsinv = lbc (gsinv, cd_type='V', nperio=nperio, psgn=-1.) 1151 #gcosv = lbc (gcosv, cd_type='V', nperio=nperio, psgn=-1.) 1152 #gsinf = lbc (gsinf, cd_type='F', nperio=nperio, psgn=-1.) 1153 #gcosf = lbc (gcosf, cd_type='F', nperio=nperio, psgn=-1.) 1154 1155 gsint = lbc_del (gsint, cd_type='T', nperio=nperio, psgn=-1.) 1156 gcost = lbc_del (gcost, cd_type='T', nperio=nperio, psgn=-1.) 1157 gsinu = lbc_del (gsinu, cd_type='U', nperio=nperio, psgn=-1.) 1158 gcosu = lbc_del (gcosu, cd_type='U', nperio=nperio, psgn=-1.) 1159 gsinv = lbc_del (gsinv, cd_type='V', nperio=nperio, psgn=-1.) 1160 gcosv = lbc_del (gcosv, cd_type='V', nperio=nperio, psgn=-1.) 1161 gsinf = lbc_del (gsinf, cd_type='F', nperio=nperio, psgn=-1.) 1162 gcosf = lbc_del (gcosf, cd_type='F', nperio=nperio, psgn=-1.) 1163 1164 if mmath == xr : 1047 1165 gsint = gsint.assign_coords ( glamt.coords ) 1048 1166 gcost = gcost.assign_coords ( glamt.coords ) … … 1058 1176 def angle (glam, gphi, nperio, cd_type='T') : 1059 1177 '''Compute sinus and cosinus of model line direction with respect to east''' 1060 math = __math__ (glam) 1178 mmath = __mmath__ (glam) 1179 1180 zlam = lbc_add (glam, nperio, cd_type, 1.) 1181 zphi = lbc_add (gphi, nperio, cd_type, 1.) 1182 1061 1183 # north pole direction & modulous 1062 zxnp = 0. - 2.0 * np.cos (rad* glam) * np.tan (rpi/4.0 - rad*gphi/2.0)1063 zynp = 0. - 2.0 * np.sin (rad* glam) * np.tan (rpi/4.0 - rad*gphi/2.0)1184 zxnp = 0. - 2.0 * np.cos (rad*zlam) * np.tan (rpi/4.0 - rad*zphi/2.0) 1185 zynp = 0. - 2.0 * np.sin (rad*zlam) * np.tan (rpi/4.0 - rad*zphi/2.0) 1064 1186 znnp = zxnp*zxnp + zynp*zynp 1065 1187 1066 1188 # j-direction: segment direction (around point) 1067 zlan_n = np.roll ( glam, axis=-2, shift=-1) # glam [jj+1, ji]1068 zphh_n = np.roll ( gphi, axis=-2, shift=-1) # gphi [jj+1, ji]1069 zlan_s = np.roll ( glam, axis=-2, shift= 1) # glam [jj-1, ji]1070 zphh_s = np.roll ( gphi, axis=-2, shift= 1) # gphi [jj-1, ji]1189 zlan_n = np.roll (zlam, axis=-2, shift=-1) # glam [jj+1, ji] 1190 zphh_n = np.roll (zphi, axis=-2, shift=-1) # gphi [jj+1, ji] 1191 zlan_s = np.roll (zlam, axis=-2, shift= 1) # glam [jj-1, ji] 1192 zphh_s = np.roll (zphi, axis=-2, shift= 1) # gphi [jj-1, ji] 1071 1193 1072 1194 zxff = 2.0 * np.cos (rad*zlan_n) * np.tan (rpi/4.0 - rad*zphh_n/2.0) \ … … 1079 1201 gcos = (zxnp*zxff + zynp*zyff) / znff 1080 1202 1081 gsin = lbc (gsin, cd_type=cd_type, nperio=nperio, psgn=-1.)1082 gcos = lbc (gcos, cd_type=cd_type, nperio=nperio, psgn=-1.)1083 1084 if m ath == xr :1203 gsin = lbc_del (gsin, cd_type=cd_type, nperio=nperio, psgn=-1.) 1204 gcos = lbc_del (gcos, cd_type=cd_type, nperio=nperio, psgn=-1.) 1205 1206 if mmath == xr : 1085 1207 gsin = gsin.assign_coords ( glam.coords ) 1086 1208 gcos = gcos.assign_coords ( glam.coords ) … … 1124 1246 east-north components on the T grid point 1125 1247 ''' 1126 m ath = __math__ (uo)1248 mmath = __mmath__ (uo) 1127 1249 1128 1250 ut = U2T (uo, nperio=nperio, psgn=-1.0, zdim=zdim) … … 1139 1261 def rot_uv2enF ( uo, vo, gsinf, gcosf, nperio, zdim='deptht' ) : 1140 1262 ''' 1141 ** Purpose : 1263 ** Purpose : Rotate the Repere: Change vector componantes from 1142 1264 stretched coordinates grid --> geographic grid 1143 1265 uo is on the U grid point, vo is on the V grid point 1144 1266 east-north components on the T grid point 1145 1267 ''' 1146 m ath = __math__ (uo)1268 mmath = __mmath__ (uo) 1147 1269 1148 1270 uf = U2F (uo, nperio=nperio, psgn=-1.0, zdim=zdim) … … 1156 1278 1157 1279 return u_e, v_n 1158 1159 def U2T (utab, nperio=None, psgn=-1.0, zdim='deptht') : 1280 1281 #@numba.jit(forceobj=True) 1282 def U2T (utab, nperio=None, psgn=-1.0, zdim='deptht', action='ave') : 1160 1283 '''Interpolate an array from U grid to T grid i-mean)''' 1161 math = __math__ (utab) 1162 utab_0 = math.where ( np.isnan(utab), 0., utab) 1163 utab_0 = lbc (utab_0 , nperio=nperio, cd_type='U', psgn=psgn) 1284 mmath = __mmath__ (utab) 1285 utab_0 = mmath.where ( np.isnan(utab), 0., utab) 1286 lperio, aperio = lbc_diag (nperio) 1287 utab_0 = lbc_add (utab_0, nperio=nperio, cd_type='U', psgn=psgn) 1164 1288 ix, ax = __findAxis__ (utab_0, 'x') 1165 1289 iz, az = __findAxis__ (utab_0, 'z') 1166 ttab = lbc ( 0.5 * (utab_0 + np.roll (utab_0, axis=ix, shift=1)), cd_type='T', nperio=nperio, psgn=psgn) 1167 1168 if math == xr : 1290 if action == 'ave' : ttab = 0.5 * (utab_0 + np.roll (utab_0, axis=ix, shift=1)) 1291 if action == 'min' : ttab = np.minimum (utab_0 , np.roll (utab_0, axis=ix, shift=1)) 1292 if action == 'max' : ttab = np.maximum (utab_0 , np.roll (utab_0, axis=ix, shift=1)) 1293 if action == 'mult' : ttab = utab_0 * np.roll (utab_0, axis=ix, shift=1) 1294 ttab = lbc_del (ttab, nperio=nperio, cd_type='T', psgn=psgn) 1295 1296 if mmath == xr : 1169 1297 if ax != None : 1170 1298 ttab = ttab.assign_coords({ax:np.arange (ttab.shape[ix])+1.}) … … 1173 1301 return ttab 1174 1302 1175 def V2T (vtab, nperio=None, psgn=-1.0, zdim='deptht') : 1303 #@numba.jit(forceobj=True) 1304 def V2T (vtab, nperio=None, psgn=-1.0, zdim='deptht', action='ave') : 1176 1305 '''Interpolate an array from V grid to T grid (j-mean)''' 1177 math = __math__ (vtab) 1178 vtab_0 = math.where ( np.isnan(vtab), 0., vtab) 1179 vtab_0 = lbc (vtab_0 , nperio=nperio, cd_type='V', psgn=psgn) 1306 mmath = __mmath__ (vtab) 1307 lperio, aperio = lbc_diag (nperio) 1308 vtab_0 = mmath.where ( np.isnan(vtab), 0., vtab) 1309 vtab_0 = lbc_add (vtab_0, nperio=nperio, cd_type='V', psgn=psgn) 1180 1310 iy, ay = __findAxis__ (vtab_0, 'y') 1181 1311 iz, az = __findAxis__ (vtab_0, 'z') 1182 ttab = lbc ( 0.5 * (vtab_0 + np.roll (vtab_0, axis=iy, shift=1)), cd_type='T', nperio=nperio, psgn=psgn) 1183 1184 if math == xr : 1312 if action == 'ave' : ttab = 0.5 * (vtab_0 + np.roll (vtab_0, axis=iy, shift=1)) 1313 if action == 'min' : ttab = np.minimum (vtab_0 , np.roll (vtab_0, axis=iy, shift=1)) 1314 if action == 'max' : ttab = np.maximum (vtab_0 , np.roll (vtab_0, axis=iy, shift=1)) 1315 if action == 'mult' : ttab = vtab_0 * np.roll (vtab_0, axis=iy, shift=1) 1316 ttab = lbc_del (ttab, nperio=nperio, cd_type='T', psgn=psgn) 1317 if mmath == xr : 1185 1318 if ay !=None : 1186 1319 ttab = ttab.assign_coords({ay:np.arange(ttab.shape[iy])+1.}) … … 1189 1322 return ttab 1190 1323 1191 def F2T (ftab, nperio=None, psgn=1.0, zdim='depthf') : 1324 #@numba.jit(forceobj=True) 1325 def F2T (ftab, nperio=None, psgn=1.0, zdim='depthf', action='ave') : 1192 1326 '''Interpolate an array from F grid to T grid (i- and j- means)''' 1193 math = __math__ (ftab) 1194 ftab_0 = math.where ( np.isnan(ftab), 0., ftab) 1195 ftab_0 = lbc (ftab_0 , nperio=nperio, cd_type='F', psgn=psgn) 1196 ttab = V2T(F2V(ftab_0, nperio=nperio, psgn=psgn, zdim=zdim), nperio=nperio, psgn=psgn, zdim=zdim) 1197 return ttab 1198 1199 def T2U (ttab, nperio=None, psgn=1.0, zdim='depthu') : 1327 mmath = __mmath__ (ftab) 1328 ftab_0 = mmath.where ( np.isnan(ftab), 0., ftab) 1329 ftab_0 = lbc_add (ftab_0 , nperio=nperio, cd_type='F', psgn=psgn) 1330 ttab = V2T(F2V(ftab_0, nperio=nperio, psgn=psgn, zdim=zdim, action=action), nperio=nperio, psgn=psgn, zdim=zdim, action=action) 1331 return lbc_del (ttab, nperio=nperio, cd_type='T', psgn=psgn) 1332 1333 #@numba.jit(forceobj=True) 1334 def T2U (ttab, nperio=None, psgn=1.0, zdim='depthu', action='ave') : 1200 1335 '''Interpolate an array from T grid to U grid (i-mean)''' 1201 m ath = __math__ (ttab)1202 ttab_0 = m ath.where ( np.isnan(ttab), 0., ttab)1203 ttab_0 = lbc (ttab_0 , nperio=nperio, cd_type='T', psgn=psgn)1336 mmath = __mmath__ (ttab) 1337 ttab_0 = mmath.where ( np.isnan(ttab), 0., ttab) 1338 ttab_0 = lbc_add (ttab_0 , nperio=nperio, cd_type='T', psgn=psgn) 1204 1339 ix, ax = __findAxis__ (ttab_0, 'x') 1205 1340 iz, az = __findAxis__ (ttab_0, 'z') 1206 utab = lbc ( 0.5 * (ttab_0 + np.roll (ttab_0, axis=ix, shift=-1)), cd_type='U', nperio=nperio, psgn=psgn) 1207 1208 if math == xr : 1341 if action == 'ave' : utab = 0.5 * (ttab_0 + np.roll (ttab_0, axis=ix, shift=-1)) 1342 if action == 'min' : utab = np.minimum (ttab_0 , np.roll (ttab_0, axis=ix, shift=-1)) 1343 if action == 'max' : utab = np.maximum (ttab_0 , np.roll (ttab_0, axis=ix, shift=-1)) 1344 if action == 'mult' : utab = ttab_0 * np.roll (ttab_0, axis=ix, shift=-1) 1345 utab = lbc_del (utab, nperio=nperio, cd_type='U', psgn=psgn) 1346 1347 if mmath == xr : 1209 1348 if ax != None : 1210 1349 utab = ttab.assign_coords({ax:np.arange(utab.shape[ix])+1.}) … … 1213 1352 return utab 1214 1353 1215 def T2V (ttab, nperio=None, psgn=1.0, zdim='depthv') : 1354 #@numba.jit(forceobj=True) 1355 def T2V (ttab, nperio=None, psgn=1.0, zdim='depthv', action='ave') : 1216 1356 '''Interpolate an array from T grid to V grid (j-mean)''' 1217 m ath = __math__ (ttab)1218 ttab_0 = m ath.where ( np.isnan(ttab), 0., ttab)1219 ttab_0 = lbc (ttab_0 , nperio=nperio, cd_type='T', psgn=psgn)1357 mmath = __mmath__ (ttab) 1358 ttab_0 = mmath.where ( np.isnan(ttab), 0., ttab) 1359 ttab_0 = lbc_add (ttab_0 , nperio=nperio, cd_type='T', psgn=psgn) 1220 1360 iy, ay = __findAxis__ (ttab_0, 'y') 1221 1361 iz, az = __findAxis__ (ttab_0, 'z') 1222 vtab = lbc ( 0.5 * (ttab_0 + np.roll (ttab_0, axis=iy, shift=-1)), cd_type='V', nperio=nperio, psgn=psgn) 1223 1224 if math == xr : 1362 if action == 'ave' : vtab = 0.5 * (ttab_0 + np.roll (ttab_0, axis=iy, shift=-1)) 1363 if action == 'min' : vtab = np.minimum (ttab_0 , np.roll (ttab_0, axis=iy, shift=-1)) 1364 if action == 'max' : vtab = np.maximum (ttab_0 , np.roll (ttab_0, axis=iy, shift=-1)) 1365 if action == 'mult' : vtab = ttab_0 * np.roll (ttab_0, axis=iy, shift=-1) 1366 1367 vtab = lbc_del (vtab, nperio=nperio, cd_type='V', psgn=psgn) 1368 if mmath == xr : 1225 1369 if ay != None : 1226 1370 vtab = vtab.assign_coords({ay:np.arange(vtab.shape[iy])+1.}) … … 1229 1373 return vtab 1230 1374 1231 def V2F (vtab, nperio=None, psgn=-1.0, zdim='depthf') : 1375 #@numba.jit(forceobj=True) 1376 def V2F (vtab, nperio=None, psgn=-1.0, zdim='depthf', action='ave') : 1232 1377 '''Interpolate an array from V grid to F grid (i-mean)''' 1233 m ath = __math__ (vtab)1234 vtab_0 = m ath.where ( np.isnan(vtab), 0., vtab)1235 vtab_0 = lbc (vtab_0 , nperio=nperio, cd_type='V', psgn=psgn)1378 mmath = __mmath__ (vtab) 1379 vtab_0 = mmath.where ( np.isnan(vtab), 0., vtab) 1380 vtab_0 = lbc_add (vtab_0 , nperio=nperio, cd_type='V', psgn=psgn) 1236 1381 ix, ax = __findAxis__ (vtab_0, 'x') 1237 1382 iz, az = __findAxis__ (vtab_0, 'z') 1238 ftab = lbc ( 0.5 * (vtab_0 + np.roll (vtab_0, axis=ix, shift=-1)), cd_type='F', nperio=nperio, psgn=psgn) 1239 1240 if math == xr : 1383 if action == 'ave' : 0.5 * (vtab_0 + np.roll (vtab_0, axis=ix, shift=-1)) 1384 if action == 'min' : np.minimum (vtab_0 , np.roll (vtab_0, axis=ix, shift=-1)) 1385 if action == 'max' : np.maximum (vtab_0 , np.roll (vtab_0, axis=ix, shift=-1)) 1386 if action == 'mult' : vtab_0 * np.roll (vtab_0, axis=ix, shift=-1) 1387 ftab = lbc_del (ftab, nperio=nperio, cd_type='F', psgn=psgn) 1388 1389 if mmath == xr : 1241 1390 if ax != None : 1242 1391 ftab = ftab.assign_coords({ax:np.arange(ftab.shape[ix])+1.}) 1243 1392 if zdim != None and iz != None and az != 'olevel' : 1244 1393 ftab = ftab.rename( {az:zdim}) 1245 return ftab 1246 1247 def U2F (utab, nperio=None, psgn=-1.0, zdim='depthf') : 1394 return lbc_del (ftab, nperio=nperio, cd_type='F', psgn=psgn) 1395 1396 #@numba.jit(forceobj=True) 1397 def U2F (utab, nperio=None, psgn=-1.0, zdim='depthf', action='ave') : 1248 1398 '''Interpolate an array from U grid to F grid i-mean)''' 1249 m ath = __math__ (utab)1250 utab_0 = m ath.where ( np.isnan(utab), 0., utab)1251 utab_0 = lbc (utab_0 , nperio=nperio, cd_type='U', psgn=psgn)1399 mmath = __mmath__ (utab) 1400 utab_0 = mmath.where ( np.isnan(utab), 0., utab) 1401 utab_0 = lbc_add (utab_0 , nperio=nperio, cd_type='U', psgn=psgn) 1252 1402 iy, ay = __findAxis__ (utab_0, 'y') 1253 1403 iz, az = __findAxis__ (utab_0, 'z') 1254 ftab = lbc ( 0.5 * (utab_0 + np.roll (utab_0, axis=iy, shift=-1)), cd_type='F', nperio=nperio, psgn=psgn) 1255 1256 if math == xr : 1404 if action == 'ave' : ftab = 0.5 * (utab_0 + np.roll (utab_0, axis=iy, shift=-1)) 1405 if action == 'min' : ftab = np.minimum (utab_0 , np.roll (utab_0, axis=iy, shift=-1)) 1406 if action == 'max' : ftab = np.maximum (utab_0 , np.roll (utab_0, axis=iy, shift=-1)) 1407 if action == 'mult' : ftab = utab_0 * np.roll (utab_0, axis=iy, shift=-1) 1408 ftab = lbc_del (ftab, nperio=nperio, cd_type='F', psgn=psgn) 1409 1410 if mmath == xr : 1257 1411 if ay != None : 1258 1412 ftab = ftab.assign_coords({'y':np.arange(ftab.shape[iy])+1.}) … … 1261 1415 return ftab 1262 1416 1263 def F2T (ftab, nperio=None, psgn=1.0, zdim='deptht') : 1417 #@numba.jit(forceobj=True) 1418 def F2T (ftab, nperio=None, psgn=1.0, zdim='deptht', action='ave') : 1264 1419 '''Interpolate an array on F grid to T grid (i- and j- means)''' 1265 math = __math__ (ftab) 1266 ftab_0 = math.where ( np.isnan(ttab), 0., ttab) 1267 ftab_0 = lbc (ftab_0 , nperio=nperio, cd_type='F', psgn=psgn) 1268 ttab = U2T(F2U(ftab_0, nperio=nperio, psgn=psgn, zdim=zdim), nperio=nperio, psgn=psgn, zdim=zdim) 1269 return ttab 1270 1271 def T2F (ttab, nperio=None, psgn=1.0, zdim='deptht') : 1420 mmath = __mmath__ (ftab) 1421 ftab_0 = mmath.where ( np.isnan(ttab), 0., ttab) 1422 ftab_0 = lbc_add (ftab_0 , nperio=nperio, cd_type='F', psgn=psgn) 1423 ttab = U2T(F2U(ftab_0, nperio=nperio, psgn=psgn, zdim=zdim, action=action), nperio=nperio, psgn=psgn, zdim=zdim, action=action) 1424 return lbc_del (ttab, nperio=nperio, cd_type='T', psgn=psgn) 1425 1426 #@numba.jit(forceobj=True) 1427 def T2F (ttab, nperio=None, psgn=1.0, zdim='deptht', action='mean') : 1272 1428 '''Interpolate an array on T grid to F grid (i- and j- means)''' 1273 math = __math__ (ftab) 1274 ttab_0 = math.where ( np.isnan(ttab), 0., ttab) 1275 ttab_0 = lbc (ttab_0 , nperio=nperio, cd_type='T', psgn=psgn) 1276 ftab = T2U(U2F(ttab, nperio=nperio, psgn=psgn, zdim=zdim), nperio=nperio, psgn=psgn, zdim=zdim) 1277 return ftab 1278 1279 def F2U (ftab, nperio=None, psgn=1.0, zdim='depthu') : 1429 mmath = __mmath__ (ttab) 1430 ttab_0 = mmath.where ( np.isnan(ttab), 0., ttab) 1431 ttab_0 = lbc_add (ttab_0 , nperio=nperio, cd_type='T', psgn=psgn) 1432 ftab = T2U(U2F(ttab, nperio=nperio, psgn=psgn, zdim=zdim, action=action), nperio=nperio, psgn=psgn, zdim=zdim, action=action) 1433 1434 return lbc_del (ftab, nperio=nperio, cd_type='F', psgn=psgn) 1435 1436 #@numba.jit(forceobj=True) 1437 def F2U (ftab, nperio=None, psgn=1.0, zdim='depthu', action='ave') : 1280 1438 '''Interpolate an array on F grid to FUgrid (i-mean)''' 1281 m ath = __math__ (ftab)1282 ftab_0 = m ath.where ( np.isnan(ftab), 0., ftab)1283 ftab_0 = lbc (ftab_0 , nperio=nperio, cd_type='F', psgn=psgn)1439 mmath = __mmath__ (ftab) 1440 ftab_0 = mmath.where ( np.isnan(ftab), 0., ftab) 1441 ftab_0 = lbc_add (ftab_0 , nperio=nperio, cd_type='F', psgn=psgn) 1284 1442 iy, ay = __findAxis__ (ftab_0, 'y') 1285 1443 iz, az = __findAxis__ (ftab_0, 'z') 1286 utab = lbc ( 0.5 * (ftab_0 + np.roll (ftab_0, axis=iy, shift=-1)), cd_type='U', nperio=nperio, psgn=psgn) 1287 1288 if math == xr : 1444 if action == 'ave' : utab = 0.5 * (ftab_0 + np.roll (ftab_0, axis=iy, shift=-1)) 1445 if action == 'min' : utab = np.minimum (ftab_0 , np.roll (ftab_0, axis=iy, shift=-1)) 1446 if action == 'max' : utab = np.maximum (ftab_0 , np.roll (ftab_0, axis=iy, shift=-1)) 1447 if action == 'mult' : utab = ftab_0 * np.roll (ftab_0, axis=iy, shift=-1) 1448 1449 utab = lbc_del (utab, nperio=nperio, cd_type='U', psgn=psgn) 1450 1451 if mmath == xr : 1289 1452 utab = utab.assign_coords({ay:np.arange(ftab.shape[iy])+1.}) 1290 1453 if zdim != None and iz != None and az != 'olevel' : … … 1292 1455 return utab 1293 1456 1294 def F2V (ftab, nperio=None, psgn=1.0, zdim='depthv') : 1457 #@numba.jit(forceobj=True) 1458 def F2V (ftab, nperio=None, psgn=1.0, zdim='depthv', action='ave') : 1295 1459 '''Interpolate an array from F grid to V grid (i-mean)''' 1296 m ath = __math__ (ftab)1297 ftab_0 = m ath.where ( np.isnan(ftab), 0., ftab)1298 ftab_0 = lbc (ftab_0 , nperio=nperio, cd_type='F', psgn=psgn)1460 mmath = __mmath__ (ftab) 1461 ftab_0 = mmath.where ( np.isnan(ftab), 0., ftab) 1462 ftab_0 = lbc_add (ftab_0 , nperio=nperio, cd_type='F', psgn=psgn) 1299 1463 ix, ax = __findAxis__ (ftab_0, 'x') 1300 1464 iz, az = __findAxis__ (ftab_0, 'z') 1301 vtab = lbc ( 0.5 * (ftab_0 + np.roll (ftab_0, axis=ix, shift=-1)), cd_type='V', nperio=nperio) 1302 1303 if math == xr : 1465 if action == 'ave' : vtab = 0.5 * (ftab_0 + np.roll (ftab_0, axis=ix, shift=-1)) 1466 if action == 'min' : vtab = np.minimum (ftab_0 , np.roll (ftab_0, axis=ix, shift=-1)) 1467 if action == 'max' : vtab = np.maximum (ftab_0 , np.roll (ftab_0, axis=ix, shift=-1)) 1468 if action == 'mult' : vtab = ftab_0 * np.roll (ftab_0, axis=ix, shift=-1) 1469 1470 vtab = lbc_del (vtab, nperio=nperio, cd_type='V', psgn=psgn) 1471 if mmath == xr : 1304 1472 vtab = vtab.assign_coords({ax:np.arange(ftab.shape[ix])+1.}) 1305 1473 if zdim != None and iz != None and az != 'olevel' : … … 1307 1475 return vtab 1308 1476 1477 #@numba.jit(forceobj=True) 1309 1478 def W2T (wtab, zcoord=None, zdim='deptht', sval=np.nan) : 1310 1479 ''' … … 1312 1481 sval is the bottom value 1313 1482 ''' 1314 m ath = __math__ (wtab)1315 wtab_0 = m ath.where ( np.isnan(wtab), 0., wtab)1483 mmath = __mmath__ (wtab) 1484 wtab_0 = mmath.where ( np.isnan(wtab), 0., wtab) 1316 1485 1317 1486 iz, az = __findAxis__ (wtab_0, 'z') … … 1319 1488 ttab = 0.5 * ( wtab_0 + np.roll (wtab_0, axis=iz, shift=-1) ) 1320 1489 1321 if m ath == xr :1490 if mmath == xr : 1322 1491 ttab[{az:iz}] = sval 1323 1492 if zdim != None and iz != None and az != 'olevel' : … … 1330 1499 return ttab 1331 1500 1501 #@numba.jit(forceobj=True) 1332 1502 def T2W (ttab, zcoord=None, zdim='depthw', sval=np.nan, extrap_surf=False) : 1333 1503 '''Interpolate an array from T grid to W grid (k-mean) … … 1335 1505 if extrap_surf==True, surface value is taken from 1st level value. 1336 1506 ''' 1337 m ath = __math__ (ttab)1338 ttab_0 = m ath.where ( np.isnan(ttab), 0., ttab)1507 mmath = __mmath__ (ttab) 1508 ttab_0 = mmath.where ( np.isnan(ttab), 0., ttab) 1339 1509 iz, az = __findAxis__ (ttab_0, 'z') 1340 1510 wtab = 0.5 * ( ttab_0 + np.roll (ttab_0, axis=iz, shift=1) ) 1341 1511 1342 if m ath == xr :1512 if mmath == xr : 1343 1513 if extrap_surf : wtab[{az:0}] = ttabb[{az:0}] 1344 1514 else : wtab[{az:0}] = sval … … 1347 1517 else : wtab[..., 0, :, :] = sval 1348 1518 1349 if m ath == xr :1519 if mmath == xr : 1350 1520 if zdim != None and iz != None and az != 'olevel' : 1351 1521 wtab = wtab.rename ( {az:zdim}) … … 1354 1524 return wtab 1355 1525 1356 def fill (ptab, nperio, cd_type='T', npass=1) : 1357 ''' 1358 Fill 0. or np.nan values with mean of neighbours 1526 #@numba.jit(forceobj=True) 1527 def fill (ptab, nperio, cd_type='T', npass=1, sval=0.) : 1528 ''' 1529 Fill sval values with mean of neighbours 1359 1530 1360 1531 Inputs : … … 1363 1534 ''' 1364 1535 1365 math = __math__ (ptab) 1366 ztab = math.where (np.isnan(ptab), 0., ptab) 1367 1368 DoPerio = False 1369 if nperio == 4.2 or nperio == 6.2 : DoPerio = True 1370 1371 if DoPerio : ztab = lbc_add (ztab) 1536 mmath = __mmath__ (ptab) 1537 1538 DoPerio = False ; lperio = nperio 1539 if nperio == 4.2 : 1540 DoPerio = True ; lperio = 4 1541 if nperio == 6.2 : 1542 DoPerio = True ; lperio = 6 1543 1544 if DoPerio : 1545 ztab = lbc_add (ptab, nperio=nperio, sval=sval) 1546 else : 1547 ztab = ptab 1548 1549 if np.isnan (sval) : 1550 ztab = mmath.where (np.isnan(ztab), np.nan, ztab) 1551 else : 1552 ztab = mmath.where (ztab==sval , np.nan, ztab) 1372 1553 1373 1554 for nn in np.arange (npass) : 1374 zmask = math.where (ztab==0., 0., 1. ) 1555 zmask = mmath.where ( np.isnan(ztab), 0., 1. ) 1556 ztab0 = mmath.where ( np.isnan(ztab), 0., ztab ) 1375 1557 # Compte du nombre de voisins 1376 zcount = zmask \1558 zcount = 1./6. * ( zmask \ 1377 1559 + np.roll(zmask, shift=1, axis=-1) + np.roll(zmask, shift=-1, axis=-1) \ 1378 1560 + np.roll(zmask, shift=1, axis=-2) + np.roll(zmask, shift=-1, axis=-2) \ … … 1381 1563 + np.roll(np.roll(zmask, shift=-1, axis=-2), shift= 1, axis=-1) \ 1382 1564 + np.roll(np.roll(zmask, shift= 1, axis=-2), shift=-1, axis=-1) \ 1383 + np.roll(np.roll(zmask, shift=-1, axis=-2), shift=-1, axis=-1) ) 1384 1385 znew = zmask \ 1386 + np.roll(ztab, shift=1, axis=-1) + np.roll(ztab, shift=-1, axis=-1) \ 1387 + np.roll(ztab, shift=1, axis=-2) + np.roll(ztab, shift=-1, axis=-2) \ 1388 + 0.5 * ( \ 1389 + np.roll(np.roll(ztab, shift= 1, axis=-2), shift= 1, axis=-1) \ 1390 + np.roll(np.roll(ztab, shift=-1, axis=-2), shift= 1, axis=-1) \ 1391 + np.roll(np.roll(ztab, shift= 1, axis=-2), shift=-1, axis=-1) \ 1392 + np.roll(np.roll(ztab, shift=-1, axis=-2), shift=-1, axis=-1) ) 1393 1394 zcount = lbc (zcount, nperio=nperio, cd_type=cd_type) 1395 znew = lbc (znew , nperio=nperio, cd_type=cd_type) 1396 1397 ztab = math.where (np.logical_and (zmask==0., zcount>0), znew/zcount, ztab) 1398 1399 ztab = math.where (ztab==0., np.nan, ztab) 1400 1401 if DoPerio : ztab = lbc_del (ztab) 1565 + np.roll(np.roll(zmask, shift=-1, axis=-2), shift=-1, axis=-1) ) ) 1566 1567 znew =1./6. * ( ztab0 \ 1568 + np.roll(ztab0, shift=1, axis=-1) + np.roll(ztab0, shift=-1, axis=-1) \ 1569 + np.roll(ztab0, shift=1, axis=-2) + np.roll(ztab0, shift=-1, axis=-2) \ 1570 + 0.5 * ( \ 1571 + np.roll(np.roll(ztab0 , shift= 1, axis=-2), shift= 1, axis=-1) \ 1572 + np.roll(np.roll(ztab0 , shift=-1, axis=-2), shift= 1, axis=-1) \ 1573 + np.roll(np.roll(ztab0 , shift= 1, axis=-2), shift=-1, axis=-1) \ 1574 + np.roll(np.roll(ztab0 , shift=-1, axis=-2), shift=-1, axis=-1) ) ) 1575 1576 zcount = lbc (zcount, nperio=lperio, cd_type=cd_type) 1577 znew = lbc (znew , nperio=lperio, cd_type=cd_type) 1578 1579 ztab = mmath.where (np.logical_and (zmask==0., zcount>0), znew/zcount, ztab) 1580 1581 ztab = mmath.where (zcount==0, sval, ztab) 1582 if DoPerio : ztab = lbc_del (ztab, nperio=lperio) 1402 1583 1403 1584 return ztab 1404 1585 1586 #@numba.jit(forceobj=True) 1405 1587 def correct_uv (u, v, lat) : 1406 1588 ''' … … 1417 1599 1418 1600 Outputs : 1419 modified eastward/nothward components have correct polar projection in cartopy 1420 ''' 1421 math = __math__ (u) 1601 modified eastward/nothward components to have correct polar projections in cartopy 1602 ''' 1422 1603 uv = np.sqrt (u*u + v*v) # Original modulus 1423 1604 zu = u … … 1427 1608 return uc, vc 1428 1609 1610 def msf (v_e1v_e3v, lat1d, depthw) : 1611 ''' 1612 Computes the meridonal stream function 1613 First input is meridional_velocity*e1v*e3v 1614 ''' 1615 @numba.jit(forceobj=True) 1616 def iin (tab, dim) : 1617 ''' 1618 Integrate from the bottom 1619 ''' 1620 result = tab * 0.0 1621 nlen = len(tab.coords[dim]) 1622 for jn in np.arange (nlen-2, 0, -1) : 1623 result [{dim:jn}] = result [{dim:jn+1}] - tab [{dim:jn}] 1624 result = result.where (result !=0, np.nan) 1625 return result 1626 1627 zomsf = iin ((v_e1v_e3v).sum (dim='x', keep_attrs=True)*1E-6, dim='depthv') 1628 zomsf = zomsf.assign_coords ( {'depthv':depthw.values, 'y':lat1d}) 1629 zomsf = zomsf.rename ( {'depthv':'depthw', 'y':'lat'}) 1630 zomsf.attrs['long_name'] = 'Meridional stream function' 1631 1632 zomsf.attrs['units'] = 'Sv' 1633 zomsf.depthw.attrs=depthw.attrs 1634 zomsf.lat.attrs=lat1d.attrs 1635 1636 return zomsf 1637 1638 def bsf (u_e2u_e3u, mask, nperio=None, bsf0=None ) : 1639 ''' 1640 Computes the barotropic stream function 1641 First input is zonal_velocity*e2u*e3u 1642 bsf0 is the point with bsf=0 (ex: bsf0={'x':5, 'y':120} ) 1643 ''' 1644 @numba.jit(forceobj=True) 1645 def iin (tab, dim) : 1646 ''' 1647 Integrate from the south 1648 ''' 1649 result = tab * 0.0 1650 nlen = len(tab.coords[dim]) 1651 for jn in np.arange (3, nlen) : 1652 result [{dim:jn}] = result [{dim:jn-1}] + tab [{dim:jn}] 1653 return result 1654 1655 bsf = iin ((u_e2u_e3u).sum(dim='depthu', keep_attrs=True)*1E-6, dim='y') 1656 bsf.attrs = u_e2u_e3u.attrs 1657 if bsf0 != None : 1658 bsf = bsf - bsf.isel (bsf0) 1659 1660 bsf = bsf.where (mask !=0, np.nan) 1661 bsf.attrs['long_name'] = 'Barotropic stream function' 1662 bsf.attrs['units'] = 'Sv' 1663 bsf = lbc (bsf, nperio=nperio, cd_type='F') 1664 1665 return bsf 1666 1667 def namelist_read (ref=None, cfg=None, out='dict', flat=False, verbose=False) : 1668 ''' 1669 Read NEMO namelist(s) and return either a dictionnary or an xarray dataset 1670 1671 ref : file with reference namelist, or a f90nml.namelist.Namelist object 1672 cfg : file with config namelist, or a f90nml.namelist.Namelist object 1673 At least one namelist neaded 1674 1675 out: 1676 'dict' to return a dictonnary 1677 'xr' to return an xarray dataset 1678 flat : only for dict output. Output a flat dictionnary with all values. 1679 1680 ''' 1681 1682 if ref != None : 1683 if isinstance (ref, str) : nml_ref = f90nml.read (ref) 1684 if isinstance (ref, f90nml.namelist.Namelist) : nml_ref = ref 1685 1686 if cfg != None : 1687 if isinstance (cfg, str) : nml_cfg = f90nml.read (cfg) 1688 if isinstance (cfg, f90nml.namelist.Namelist) : nml_cfg = cfg 1689 1690 if out == 'dict' : dict_namelist = {} 1691 if out == 'xr' : xr_namelist = xr.Dataset () 1692 1693 list_nml = [] ; list_comment = [] 1694 1695 if ref != None : 1696 list_nml.append (nml_ref) ; list_comment.append ('ref') 1697 if cfg != None : 1698 list_nml.append (nml_cfg) ; list_comment.append ('cfg') 1699 1700 for nml, comment in zip (list_nml, list_comment) : 1701 if verbose : print (comment) 1702 if flat and out =='dict' : 1703 for nam in nml.keys () : 1704 if verbose : print (nam) 1705 for value in nml[nam] : 1706 if out == 'dict' : dict_namelist[value] = nml[nam][value] 1707 if verbose : print (nam, ':', value, ':', nml[nam][value]) 1708 else : 1709 for nam in nml.keys () : 1710 if verbose : print (nam) 1711 if out == 'dict' : 1712 if nam not in dict_namelist.keys () : dict_namelist[nam] = {} 1713 for value in nml[nam] : 1714 if out == 'dict' : dict_namelist[nam][value] = nml[nam][value] 1715 if out == 'xr' : xr_namelist[value] = nml[nam][value] 1716 if verbose : print (nam, ':', value, ':', nml[nam][value]) 1717 1718 if out == 'dict' : return dict_namelist 1719 if out == 'xr' : return xr_namelist 1720 1721 1722 def fill_closed_seas (imask, nperio=None, cd_type='T') : 1723 '''Fill closed seas with image processing library 1724 imask : mask, 1 on ocean, 0 on land 1725 ''' 1726 from scipy import ndimage 1727 1728 imask_filled = ndimage.binary_fill_holes ( lbc (imask, nperio=nperio, cd_type=cd_type)) 1729 imask_filled = lbc ( imask_filled, nperio=nperio, cd_type=cd_type) 1730 1731 return imask_filled 1732 1429 1733 ## =========================================================================== 1430 1734 ## … … 1433 1737 ## =========================================================================== 1434 1738 1435 def is_orca_north_fold( Xtest, cname_long='T' ) :1739 def __is_orca_north_fold__ ( Xtest, cname_long='T' ) : 1436 1740 ''' 1437 1741 Ported (pirated !!?) from Sosie
Note: See TracChangeset
for help on using the changeset viewer.