Ignore:
Timestamp:
10/10/23 12:58:04 (7 months ago)
Author:
omamce
Message:

O.M. :

New version of WATER_BUDGET

  • Conservation in NEMO are OK
  • Conservation in Sechiba are OK, for both ICO and latlon grids
  • Problems in atmosphere, LIC surface and ocean/atmosphere coherence
File:
1 edited

Legend:

Unmodified
Added
Removed

  • TOOLS/WATER_BUDGET/nemo.py

    r6519 r6647  
    2222Periodicity and other stuff 
    2323 
     24- Lots of tests for xarray object 
     25- Not much testerd for numpy objects 
     26 
    2427olivier.marti@lsce.ipsl.fr 
    25 ''' 
    2628 
    2729## SVN information 
    … …  
    2931__Date__     = "$Date$" 
    3032__Revision__ = "$Revision$" 
    31 __Id__       = "$Id: nemo.py 6265 2022-11-02 12:45:56Z omamce $" 
     33__Id__       = "$Id: $" 
    3234__HeadURL    = "$HeadURL$" 
     35''' 
    3336 
    3437import numpy as np 
    … …  
    3639except ImportError : pass 
    3740 
    38 try    : import f90nml 
    39 except : pass 
    40  
    41 try : from sklearn.impute import SimpleImputer 
    42 except : pass 
    43  
    44 try    : import numba 
    45 except : pass 
     41#try    : import f90nml 
     42#except : pass 
     43 
     44#try : from sklearn.impute import SimpleImputer 
     45#except : pass 
    4646 
    4747rpi = np.pi ; rad = np.deg2rad (1.0) ; dar = np.rad2deg (1.0) 
    … …  
    6161repsi  = np.finfo (1.0).eps 
    6262 
    63 xList = [ 'x', 'X', 'lon'   , 'longitude' ] 
    64 yList = [ 'y', 'Y', 'lat'   , 'latitude'  ] 
    65 zList = [ 'z', 'Z', 'depth' , ] 
    66 tList = [ 't', 'T', 'time'  , ] 
     63## Default names of dimensions 
     64dim_names = {'x':'xx', 'y':'yy', 'z':'olevel', 't':None} 
     65 
     66## All possibles name of dimensions in Nemo files 
     67xName = [ 'x', 'X', 'X1', 'xx', 'XX', 'x_grid_T', 'x_grid_U', 'x_grid_V', 'x_grid_F', 'x_grid_W', 'lon', 'nav_lon', 'longitude', 'X1', 'x_c', 'x_f', ] 
     68yName = [ 'y', 'Y', 'Y1', 'yy', 'YY', 'y_grid_T', 'y_grid_U', 'y_grid_V', 'y_grid_F', 'y_grid_W', 'lat', 'nav_lat', 'latitude' , 'Y1', 'y_c', 'y_f', ] 
     69zName = [ 'z', 'Z', 'Z1', 'zz', 'ZZ', 'depth', 'tdepth', 'udepth', 'vdepth', 'wdepth', 'fdepth', 'deptht', 'depthu', 'depthv', 'depthw', 'depthf', 'olevel', 'z_c', 'z_f', ] 
     70tName = [ 't', 'T', 'tt', 'TT', 'time', 'time_counter', 'time_centered', ] 
     71 
     72## All possibles name of units of dimensions in Nemo files 
     73xUnit = [ 'degrees_east', ] 
     74yUnit = [ 'degrees_north', ] 
     75zUnit = [ 'm', 'meter', ] 
     76tUnit = [ 'second', 'minute', 'hour', 'day', 'month', 'year', ] 
     77 
     78## All possibles size of dimensions in Orca files 
     79xLength = [ 180, 182, 360, 362 ] 
     80yLength = [ 148, 149, 331, 332 ] 
     81zLength = [31, 75] 
    6782 
    6883## =========================================================================== 
    6984def __mmath__ (tab, default=None) : 
     85    ''' 
     86    Determines the type of tab : xarray or numpy object ? 
     87    ''' 
    7088    mmath = default 
    7189    try    : 
    7290        if type (tab) == xr.core.dataarray.DataArray : mmath = xr 
    73     except : 
    74         pass 
     91    except : pass 
    7592 
    7693    try    : 
    7794        if type (tab) == np.ndarray : mmath = np 
    78     except : 
    79         pass 
     95    except : pass 
    8096             
    8197    return mmath 
    82  
    8398 
    8499def __guessNperio__ (jpj, jpi, nperio=None, out='nperio') : 
    … …  
    99114    ''' 
    100115    Tries to guess the value of nperio (periodicity parameter. See NEMO documentation for details) 
    101      
     116 
    102117    Inputs 
    103118    jpj    : number of latitudes 
    … …  
    105120    nperio : periodicity parameter 
    106121    ''' 
     122    print ( jpi, jpj) 
    107123    if nperio == None : 
    108124        ## Values for NEMO version < 4.2 
    109         if jpj ==  149 and jpi ==  182 : 
     125        if (jpj ==  149 and jpi == 182) or (jpj == None and jpi == 182) or (jpj == 149 or jpi == None) : 
    110126            config = 'ORCA2.3' 
    111127            nperio = 4   # ORCA2. We choose legacy orca2. 
    112128            Iperio = 1 ; Jperio = 0 ; NFold = 1 ; NFtype = 'T' 
    113         if jpj ==  332 and jpi ==  362 : # eORCA1. 
     129        if (jpj == 332 and jpi == 362) or (jpj == None and jpi == 362) or (jpj ==  332 and jpi == None) : # eORCA1. 
    114130            config = 'eORCA1.2' 
    115131            nperio = 6   
    … …  
    125141             
    126142        ## Values for NEMO version >= 4.2. No more halo points 
    127         if jpj == 148 and jpi ==  180 : 
     143        if (jpj == 148 and jpi == 180) or (jpj == None and jpi == 180) or (jpj == 148 and jpi == None) : 
    128144            config = 'ORCA2.4' 
    129145            nperio = 4.2 # ORCA2. We choose legacy orca2. 
    130146            Iperio = 1 ; Jperio = 0 ; NFold = 1 ; NFtype = 'F' 
    131         if jpj == 331  and jpi ==  360 : # eORCA1. 
     147        if (jpj == 331 and jpi == 360) or (jpj == None and jpi == 360) or (jpj == 331 and jpi == None) : # eORCA1. 
    132148            config = 'eORCA1.4' 
    133149            nperio = 6.2 
    … …  
    142158        else : 
    143159            if nperio in nperio_valid_range : 
    144                 print ('nperio set as {:} (deduced from jpj={:d} jpi={:d})'.format (nperio, jpj, jpi)) 
     160                print ( f'nperio set as {nperio} (deduced from {jpj=} and {jpi=})' ) 
    145161            else :  
    146                 raise ValueError ('nperio set as {:} (deduced from jpi={:d}) : nemo.py is not ready for this value'.format (nperio, jpi)) 
     162                raise ValueError ( f'nperio set as {nperio} (deduced from {jpi=} and {jpj=}) : nemo.py is not ready for this value' ) 
    147163 
    148164    if out == 'nperio' : return nperio 
    … …  
    162178    Credits : who is the original author ? 
    163179    ''' 
     180     
    164181    gP = None 
    165182    mmath = __mmath__ (ptab) 
    … …  
    178195            raise Exception ('in nemo module : cd_type not found, and cannot by guessed') 
    179196        else : 
    180             print ('Grid set as', gP, 'deduced from dims ', ptab.dims) 
     197            print ( f'Grid set as {gP} deduced from dims {ptab.dims}' ) 
    181198            return gP 
    182199    else : 
    183200         raise Exception  ('in nemo module : cd_type not found, input is not an xarray data') 
    184201 
     202def get_shape ( ptab ) : 
     203    ''' 
     204    Get shape of ptab :  
     205    shape main contain X, Y, Z or T 
     206    Y is missing for a latitudinal slice 
     207    X is missing for on longitudinal slice 
     208    etc ... 
     209    ''' 
     210     
     211    get_shape = '' 
     212    ix, ax = __findAxis__ (ptab, 'x') 
     213    jy, ay = __findAxis__ (ptab, 'y') 
     214    kz, az = __findAxis__ (ptab, 'z') 
     215    lt, at = __findAxis__ (ptab, 't') 
     216    if ax : get_shape = 'X' 
     217    if ay : get_shape = 'Y' + get_shape 
     218    if az : get_shape = 'Z' + get_shape 
     219    if at : get_shape = 'T' + get_shape 
     220    return get_shape 
     221      
    185222def lbc_diag (nperio) : 
    186223    lperio = nperio ; aperio = False 
    … …  
    190227        lperio = 6 ; aperio = True 
    191228         
    192     return lperio, aperio  
     229    return lperio, aperio 
    193230 
    194231def __findAxis__ (tab, axis='z') : 
    195232    ''' 
    196     Find number and name of the requested axis 
     233    Find order and name of the requested axis 
    197234    ''' 
    198235    mmath = __mmath__ (tab) 
    199236    ix = None ; ax = None 
    200237 
    201     if axis in xList : 
    202         axList = [ 'x', 'X', 
    203                    'lon', 'nav_lon', 'nav_lon_T', 'nav_lon_U', 'nav_lon_V', 'nav_lon_F', 'nav_lon_W', 
    204                    'x_grid_T', 'x_grid_U', 'x_grid_V', 'x_grid_F', 'x_grid_W', 
    205                    'glam', 'glamt', 'glamu', 'glamv', 'glamf', 'glamw' ] 
    206         unList = [ 'degrees_east' ] 
    207     if axis in yList : 
    208         axList = [ 'y', 'Y', 'lat', 
    209                    'nav_lat', 'nav_lat_T', 'nav_lat_U', 'nav_lat_V', 'nav_lat_F', 'nav_lat_W', 
    210                    'y_grid_T', 'y_grid_U', 'y_grid_V', 'y_grid_F', 'y_grid_W', 
    211                    'gphi', 'gphi', 'gphiu', 'gphiv', 'gphif', 'gphiw'] 
    212         unList = [ 'degrees_north' ] 
    213     if axis in zList : 
    214         axList = [ 'z', 'Z', 
    215                    'depth', 'deptht', 'depthu', 'depthv', 'depthf', 'depthw', 
    216                    'olevel' ] 
    217         unList = [ 'm', 'meter' ] 
    218     if axis in tList : 
    219         axList = [ 't', 'T', 'time', 'time_counter' ] 
    220         unList = [ 'second', 'minute', 'hour', 'day', 'month' ] 
     238    if axis in xName : axName = xName ; unList = xUnit ; Length = xLength 
     239    if axis in yName : axName = yName ; unList = yUnit ; Length = yLength 
     240    if axis in zName : axName = zName ; unList = zUnit ; Length = zLength 
     241    if axis in tName : axName = tName ; unList = tUnit ; Length = None 
    221242     
    222243    if mmath == xr : 
    223         for Name in axList : 
     244        for Name in axName : 
    224245            try    : 
    225246                ix = tab.dims.index (Name) 
    … …  
    231252                for name in unList : 
    232253                    if name in tab.coords[dim].attrs['units'] : 
    233                         ix = i 
    234                         ax = dim 
     254                        ix = i ; ax = dim 
    235255    else : 
    236         if axis in xList : ix=-1 
    237         if axis in yList : 
    238             if len(tab.shape) >= 2 : ix=-2 
    239         if axis in zList : 
    240             if len(tab.shape) >= 3 : ix=-3 
    241         if axis in tList : 
    242             if len(tab.shape) >=3  : ix=-3 
    243             if len(tab.shape) >=4  : ix=-4 
     256        #if axis in xName : ix=-1 
     257        #if axis in yName : 
     258        #    if len(tab.shape) >= 2 : ix=-2 
     259        #if axis in zName : 
     260        #    if len(tab.shape) >= 3 : ix=-3 
     261        #if axis in tName : 
     262        #    if len(tab.shape) >=3  : ix=-3 
     263        #    if len(tab.shape) >=4  : ix=-4 
     264 
     265        l_shape = tab.shape 
     266        for nn in np.arange ( len(l_shape)) : 
     267            if l_shape[nn] in Length : ix = nn 
    244268        
    245269    return ix, ax 
    246270 
    247 #@numba.jit(forceobj=True) 
     271def findAxis ( tab, axis= 'z' ) : 
     272  ix, xx = __findAxis__ (tab, axis) 
     273  return xx 
     274 
    248275def fixed_lon (lon, center_lon=0.0) : 
    249276    ''' 
    … …  
    277304    return fixed_lon 
    278305 
    279 #@numba.jit(forceobj=True) 
     306def bounds_clolon ( bounds_lon, lon, rad=False, deg=True) : 
     307    '''Choose closest to lon0 longitude, adding or substacting 360° if needed''' 
     308 
     309    if rad : lon_range = 2.0*np.pi 
     310    if deg : lon_range = 360.0 
     311    bounds_clolon = bounds_lon.copy () 
     312 
     313    bounds_clolon = xr.where ( bounds_clolon < lon-lon_range/2., bounds_clolon+lon_range, bounds_clolon ) 
     314    bounds_clolon = xr.where ( bounds_clolon > lon+lon_range/2., bounds_clolon-lon_range, bounds_clolon ) 
     315 
     316    return bounds_clolon 
     317 
     318def UnifyDims ( dd, udims=dim_names, verbose=False ) : 
     319    ''' 
     320    Rename dimensions to unify them between NEMO versions 
     321    ''' 
     322     
     323    if udims['x'] : 
     324        for xx in xName : 
     325            if xx in dd.dims and xx != udims['x'] : 
     326                if verbose : print ( f"{xx} renamed to {udims['x']}" ) 
     327                dd = dd.rename ( {xx:udims['x']}) 
     328    if udims['y'] : 
     329        for yy in yName : 
     330            if yy in dd.dims and yy != udims['y']  : 
     331                if verbose : print ( f"{yy} renamed to {udims['y']}" ) 
     332                dd = dd.rename ( {yy:udims['y']} ) 
     333    if udims['z'] : 
     334        for zz in zName : 
     335            if zz in dd.dims and zz != udims['z'] : 
     336                if verbose : print ( f"{zz} renamed to {udims['z']}" ) 
     337                dd = dd.rename ( {zz:udims['z']} ) 
     338    if udims['t'] : 
     339        for tt in tName  : 
     340            if tt in dd.dims and tt != udims['t'] : 
     341                if verbose : print ( f"{tt} renamed to {udims['t']}" ) 
     342                dd = dd.rename ( {tt:udims['t']} ) 
     343 
     344    return dd 
     345 
    280346def fill_empty (ztab, sval=np.nan, transpose=False) : 
    281347    ''' 
    … …  
    283349 
    284350    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     ''' 
     351    some parts of the domain, with no ocean points, have no 
     352    values 
     353    ''' 
     354    from sklearn.impute import SimpleImputer 
    288355    mmath = __mmath__ (ztab) 
    289356 
    … …  
    302369    return ptab 
    303370 
    304 #@numba.jit(forceobj=True) 
    305371def fill_lonlat (lon, lat, sval=-1) : 
    306372    ''' 
    … …  
    308374 
    309375    Useful when NEMO has run with no wet points options :  
    310     some parts of the domain, with no ocean points, as no 
     376    some parts of the domain, with no ocean points, have no 
    311377    lon/lat values 
    312378    ''' 
     379    from sklearn.impute import SimpleImputer 
    313380    mmath = __mmath__ (lon) 
    314381 
    … …  
    328395    return plon, plat 
    329396 
    330 #@numba.jit(forceobj=True) 
     397def fill_bounds_lonlat (bounds_lon, bounds_lat, sval=-1) : 
     398    ''' 
     399    Fill longitude/latitude bounds values 
     400 
     401    Useful when NEMO has run with no wet points options :  
     402    some parts of the domain, with no ocean points, as no 
     403    lon/lat values 
     404    ''' 
     405    mmath = __mmath__ (bounds_lon) 
     406 
     407    p_bounds_lon = np.empty ( bounds_lon.shape ) 
     408    p_bounds_lat = np.empty ( bounds_lat.shape ) 
     409 
     410    imp = SimpleImputer (missing_values=sval, strategy='mean') 
     411     
     412    for n in np.arange (4) :  
     413        imp.fit (bounds_lon[:,:,n]) 
     414        p_bounds_lon[:,:,n] = imp.transform (bounds_lon[:,:,n]) 
     415        imp.fit (bounds_lat[:,:,n].T) 
     416        p_bounds_lat[:,:,n] = imp.transform (bounds_lat[:,:,n].T).T 
     417         
     418    if mmath == xr : 
     419        p_bounds_lon = xr.DataArray (bounds_lon, dims=bounds_lon.dims, coords=bounds_lon.coords) 
     420        p_bounds_lat = xr.DataArray (bounds_lat, dims=bounds_lat.dims, coords=bounds_lat.coords) 
     421        p_bounds_lon.attrs = bounds_lat.attrs ; p_bounds_lat.attrs = bounds_lat.attrs 
     422         
     423    return p_bounds_lon, p_bounds_lat 
     424 
    331425def jeq (lat) : 
    332426    ''' 
    … …  
    337431    mmath = __mmath__ (lat) 
    338432    ix, ax = __findAxis__ (lat, 'x') 
    339     iy, ay = __findAxis__ (lat, 'y') 
     433    jy, ay = __findAxis__ (lat, 'y') 
    340434 
    341435    if mmath == xr : 
    342         jeq = int ( np.mean ( np.argmin (np.abs (np.float64 (lat)), axis=iy) ) ) 
     436        jeq = int ( np.mean ( np.argmin (np.abs (np.float64 (lat)), axis=jy) ) ) 
    343437    else :  
    344438        jeq = np.argmin (np.abs (np.float64 (lat[...,:, 0]))) 
    345439    return jeq 
    346440 
    347 #@numba.jit(forceobj=True) 
    348441def lon1D (lon, lat=None) : 
    349442    ''' 
    … …  
    354447    ''' 
    355448    mmath = __mmath__ (lon) 
    356     if np.max (lat) != None : 
     449    jpj, jpi  = lon.shape [-2:] 
     450    if np.max (lat) : 
    357451        je    = jeq (lat) 
    358         lon1D = lon.copy() [..., je, :] 
     452        #lon1D = lon.copy() [..., je, :] 
     453        lon0 = lon [..., je, 0].copy() 
     454        dlon = lon [..., je, 1].copy() - lon [..., je, 0].copy() 
     455        lon1D = np.linspace ( start=lon0, stop=lon0+360.+2*dlon, num=jpi ) 
    359456    else : 
    360         jpj, jpi = lon.shape [-2:] 
    361         lon1D    = lon.copy() [..., jpj//3, :] 
    362  
    363     start = np.argmax (np.abs (np.diff (lon1D, axis=-1)) > 180.0, axis=-1) 
    364     lon1D [..., start+1:] += 360 
     457        lon0 = lon [..., jpj//3, 0].copy() 
     458        dlon = lon [..., jpj//3, 1].copy() - lon [..., jpj//3, 0].copy() 
     459        lon1D = np.linspace ( start=lon0, stop=lon0+360.+2*dlon, num=jpi ) 
     460 
     461    #start = np.argmax (np.abs (np.diff (lon1D, axis=-1)) > 180.0, axis=-1) 
     462    #lon1D [..., start+1:] += 360 
    365463 
    366464    if mmath == xr : 
     465        lon1D = xr.DataArray( lon1D, dims=('lon',), coords=(lon1D,)) 
    367466        lon1D.attrs = lon.attrs 
    368         lon1D = lon1D.assign_coords ( {'x':lon1D} ) 
     467        lon1D.attrs['units']         = 'degrees_east' 
     468        lon1D.attrs['standard_name'] = 'longitude' 
     469        lon1D.attrs['long_name :']   = 'Longitude' 
    369470         
    370471    return lon1D 
    371472 
    372 #@numba.jit(forceobj=True) 
    373473def latreg (lat, diff=0.1) : 
    374474    ''' 
    … …  
    381481    if diff == None : 
    382482        dy   = np.float64 (np.mean (np.abs (lat - np.roll (lat,shift=1,axis=-2, roll_coords=False)))) 
     483        print ( f'{dy=}' ) 
    383484        diff = dy/100. 
    384      
     485 
    385486    je     = jeq (lat) 
    386487    jreg   = np.where (lat[...,je:,:].max(axis=-1) - lat[...,je:,:].min(axis=-1)< diff)[-1][-1] + je 
    … …  
    390491    return jreg, latreg 
    391492 
    392 #@numba.jit(forceobj=True) 
    393493def lat1D (lat) : 
    394494    ''' 
    … …  
    403503    je     = jeq (lat) 
    404504    lat_eq = np.float64 (lat[...,je,:].mean(axis=-1)) 
    405        
     505      
    406506    jreg, lat_reg = latreg (lat) 
    407507    lat_ave = np.mean (lat, axis=-1) 
    408508 
     509    #print ( f'{dy=} {jpj=} {je=} {lat_eq=} {jreg=} ' ) 
     510     
    409511    if (np.abs (lat_eq) < dy/100.) : # T, U or W grid 
    410         dys    = (90.-lat_reg) / (jpj-jreg-1)*0.5 
    411         yrange = 90.-dys-lat_reg 
     512        if jpj-1 > jreg : dys = (90.-lat_reg) / (jpj-jreg-1)*0.5 
     513        else            : dys = (90.-lat_reg) / 2.0 
     514        yrange = (90.-dys-lat_reg) 
    412515    else                           :  # V or F grid 
    413         yrange = 90.    -lat_reg 
    414          
    415     lat1D = mmath.where (lat_ave<lat_reg, lat_ave, lat_reg + yrange * (np.arange(jpj)-jreg)/(jpj-jreg-1))    
    416          
     516        yrange = 90.-lat_reg 
     517 
     518    if jpj-1 > jreg : 
     519        lat1D = mmath.where (lat_ave<lat_reg, lat_ave, lat_reg + yrange * (np.arange(jpj)-jreg)/(jpj-jreg-1) ) 
     520    else : 
     521        lat1D = lat_ave 
     522    lat1D[-1] = 90.0 
     523 
    417524    if mmath == xr : 
     525        lat1D = xr.DataArray( lat1D.values, dims=('lat',), coords=(lat1D,)) 
    418526        lat1D.attrs = lat.attrs 
    419         lat1D = lat1D.assign_coords ( {'y':lat1D} ) 
    420  
     527        lat1D.attrs ['units']         = 'degrees_north' 
     528        lat1D.attrs ['standard_name'] = 'latitude' 
     529        lat1D.attrs ['long_name :']   = 'Latitude' 
     530         
    421531    return lat1D 
    422532 
    423 #@numba.jit(forceobj=True) 
    424533def latlon1D (lat, lon) : 
    425534    ''' 
    … …  
    430539    return lat1D (lat),  lon1D (lon, lat) 
    431540 
    432 #@numba.jit(forceobj=True) 
    433541def mask_lonlat (ptab, x0, x1, y0, y1, lon, lat, sval=np.nan) : 
    434542    mmath = __mmath__ (ptab) 
    … …  
    445553    return tab 
    446554 
    447 #@numba.jit(forceobj=True)       
    448555def extend (tab, Lon=False, jplus=25, jpi=None, nperio=4) : 
    449556    ''' 
    450557    Returns extended field eastward to have better plots, and box average crossing the boundary 
    451558    Works only for xarray and numpy data (?) 
     559 
     560    Useful for vertical sections in OCE and ATM. 
    452561 
    453562    tab : field to extend. 
    … …  
    539648    See NEMO documentation for further details 
    540649    ''' 
    541     jpj, jpi = ptab.shape[-2:] 
     650    jpi = None ; jpj = None 
     651    ix, ax = __findAxis__ (ptab, 'x') 
     652    jy, ay = __findAxis__ (ptab, 'y') 
     653    if ax : jpi = ptab.shape[ix] 
     654    if ay : jpj = ptab.shape[jy] 
     655         
    542656    if nperio == None : nperio = __guessNperio__ (jpj, jpi, nperio) 
    543657     
    544658    if nperio not in nperio_valid_range : 
    545         raise Exception ('nperio=', nperio, ' is not in the valid range', nperio_valid_range) 
     659        raise Exception ( f'{nperio=} is not in the valid range {nperio_valid_range}' ) 
    546660 
    547661    return jpj, jpi, nperio 
    548662         
    549 #@numba.jit(forceobj=True) 
    550663def lbc (ptab, nperio=None, cd_type='T', psgn=1.0, nemo_4U_bug=False) : 
    551664    ''' 
    … …  
    559672    ''' 
    560673    jpj, jpi, nperio = lbc_init (ptab, nperio) 
     674    ix, ax = __findAxis__ (ptab, 'x') 
     675    jy, ay = __findAxis__ (ptab, 'y') 
    561676    psgn   = ptab.dtype.type (psgn) 
    562677    mmath = __mmath__ (ptab) 
    … …  
    565680    else           : ztab = ptab.copy () 
    566681         
    567     # 
    568     #> East-West boundary conditions 
    569     # ------------------------------ 
    570     if nperio in [1, 4, 6] : 
     682    if ax :  
     683        # 
     684        #> East-West boundary conditions 
     685        # ------------------------------ 
     686        if nperio in [1, 4, 6] : 
    571687        # ... cyclic 
    572         ztab [..., :,  0] = ztab [..., :, -2] 
    573         ztab [..., :, -1] = ztab [..., :,  1] 
    574     # 
    575     #> North-South boundary conditions 
    576     # -------------------------------- 
    577     if nperio in [3, 4] :  # North fold T-point pivot 
    578         if cd_type in [ 'T', 'W' ] : # T-, W-point 
    579             ztab [..., -1, 1:       ] = psgn * ztab [..., -3, -1:0:-1      ] 
    580             ztab [..., -1, 0        ] = psgn * ztab [..., -3, 2            ] 
    581             ztab [..., -2, jpi//2:  ] = psgn * ztab [..., -2, jpi//2:0:-1  ] 
     688            ztab [...,  0] = ztab [..., -2] 
     689            ztab [..., -1] = ztab [...,  1] 
     690 
     691        if ay :  
     692            # 
     693            #> North-South boundary conditions 
     694            # -------------------------------- 
     695            if nperio in [3, 4] :  # North fold T-point pivot 
     696                if cd_type in [ 'T', 'W' ] : # T-, W-point 
     697                    ztab [..., -1, 1:       ] = psgn * ztab [..., -3, -1:0:-1      ] 
     698                    ztab [..., -1, 0        ] = psgn * ztab [..., -3, 2            ] 
     699                    ztab [..., -2, jpi//2:  ] = psgn * ztab [..., -2, jpi//2:0:-1  ] 
     700                     
     701                if cd_type == 'U' : 
     702                    ztab [..., -1, 0:-1     ] = psgn * ztab [..., -3, -1:0:-1      ]        
     703                    ztab [..., -1,  0       ] = psgn * ztab [..., -3,  1           ] 
     704                    ztab [..., -1, -1       ] = psgn * ztab [..., -3, -2           ] 
     705                     
     706                    if nemo_4U_bug : 
     707                        ztab [..., -2, jpi//2+1:-1] = psgn * ztab [..., -2, jpi//2-2:0:-1] 
     708                        ztab [..., -2, jpi//2-1   ] = psgn * ztab [..., -2, jpi//2       ] 
     709                    else : 
     710                        ztab [..., -2, jpi//2-1:-1] = psgn * ztab [..., -2, jpi//2:0:-1] 
     711                         
     712                if cd_type == 'V' :  
     713                    ztab [..., -2, 1:       ] = psgn * ztab [..., -3, jpi-1:0:-1   ] 
     714                    ztab [..., -1, 1:       ] = psgn * ztab [..., -4, -1:0:-1      ]    
     715                    ztab [..., -1, 0        ] = psgn * ztab [..., -4, 2            ] 
     716                     
     717                if cd_type == 'F' : 
     718                    ztab [..., -2, 0:-1     ] = psgn * ztab [..., -3, -1:0:-1      ] 
     719                    ztab [..., -1, 0:-1     ] = psgn * ztab [..., -4, -1:0:-1      ] 
     720                    ztab [..., -1,  0       ] = psgn * ztab [..., -4,  1           ] 
     721                    ztab [..., -1, -1       ] = psgn * ztab [..., -4, -2           ] 
    582722                 
    583         if cd_type == 'U' : 
    584             ztab [..., -1, 0:-1     ] = psgn * ztab [..., -3, -1:0:-1      ]        
    585             ztab [..., -1,  0       ] = psgn * ztab [..., -3,  1           ] 
    586             ztab [..., -1, -1       ] = psgn * ztab [..., -3, -2           ] 
     723            if nperio in [4.2] :  # North fold T-point pivot 
     724                if cd_type in [ 'T', 'W' ] : # T-, W-point 
     725                    ztab [..., -1, jpi//2:  ] = psgn * ztab [..., -1, jpi//2:0:-1  ] 
     726                     
     727                if cd_type == 'U' : 
     728                    ztab [..., -1, jpi//2-1:-1] = psgn * ztab [..., -1, jpi//2:0:-1] 
     729                     
     730                if cd_type == 'V' :  
     731                    ztab [..., -1, 1:       ] = psgn * ztab [..., -2, jpi-1:0:-1   ] 
     732                     
     733                if cd_type == 'F' : 
     734                    ztab [..., -1, 0:-1     ] = psgn * ztab [..., -2, -1:0:-1      ] 
    587735                 
    588             if nemo_4U_bug : 
    589                 ztab [..., -2, jpi//2+1:-1] = psgn * ztab [..., -2, jpi//2-2:0:-1] 
    590                 ztab [..., -2, jpi//2-1   ] = psgn * ztab [..., -2, jpi//2       ] 
    591             else : 
    592                 ztab [..., -2, jpi//2-1:-1] = psgn * ztab [..., -2, jpi//2:0:-1] 
    593                  
    594         if cd_type == 'V' :  
    595             ztab [..., -2, 1:       ] = psgn * ztab [..., -3, jpi-1:0:-1   ] 
    596             ztab [..., -1, 1:       ] = psgn * ztab [..., -4, -1:0:-1      ]    
    597             ztab [..., -1, 0        ] = psgn * ztab [..., -4, 2            ] 
    598                  
    599         if cd_type == 'F' : 
    600             ztab [..., -2, 0:-1     ] = psgn * ztab [..., -3, -1:0:-1      ] 
    601             ztab [..., -1, 0:-1     ] = psgn * ztab [..., -4, -1:0:-1      ] 
    602             ztab [..., -1,  0       ] = psgn * ztab [..., -4,  1           ] 
    603             ztab [..., -1, -1       ] = psgn * ztab [..., -4, -2           ] 
    604  
    605     if nperio in [4.2] :  # North fold T-point pivot 
    606         if cd_type in [ 'T', 'W' ] : # T-, W-point 
    607             ztab [..., -1, jpi//2:  ] = psgn * ztab [..., -1, jpi//2:0:-1  ] 
    608                  
    609         if cd_type == 'U' : 
    610             ztab [..., -1, jpi//2-1:-1] = psgn * ztab [..., -1, jpi//2:0:-1] 
    611                  
    612         if cd_type == 'V' :  
    613             ztab [..., -1, 1:       ] = psgn * ztab [..., -2, jpi-1:0:-1   ] 
    614                  
    615         if cd_type == 'F' : 
    616             ztab [..., -1, 0:-1     ] = psgn * ztab [..., -2, -1:0:-1      ] 
    617  
    618     if nperio in [5, 6] :            #  North fold F-point pivot   
    619         if cd_type in ['T', 'W']  : 
    620             ztab [..., -1, 0:       ] = psgn * ztab [..., -2, -1::-1       ] 
    621                  
    622         if cd_type == 'U' : 
    623             ztab [..., -1, 0:-1     ] = psgn * ztab [..., -2, -2::-1       ]        
    624             ztab [..., -1, -1       ] = psgn * ztab [..., -2, 0            ] # Bug ? 
    625                  
    626         if cd_type == 'V' : 
    627             ztab [..., -1, 0:       ] = psgn * ztab [..., -3, -1::-1       ] 
    628             ztab [..., -2, jpi//2:  ] = psgn * ztab [..., -2, jpi//2-1::-1 ] 
    629                  
    630         if cd_type == 'F' : 
    631             ztab [..., -1, 0:-1     ] = psgn * ztab [..., -3, -2::-1       ] 
    632             ztab [..., -1, -1       ] = psgn * ztab [..., -3, 0            ] 
    633             ztab [..., -2, jpi//2:-1] = psgn * ztab [..., -2, jpi//2-2::-1 ] 
    634  
    635     # 
    636     #> East-West boundary conditions 
    637     # ------------------------------ 
    638     if nperio in [1, 4, 6] : 
    639         # ... cyclic 
    640         ztab [..., :,  0] = ztab [..., :, -2] 
    641         ztab [..., :, -1] = ztab [..., :,  1] 
     736            if nperio in [5, 6] :            #  North fold F-point pivot   
     737                if cd_type in ['T', 'W']  : 
     738                    ztab [..., -1, 0:       ] = psgn * ztab [..., -2, -1::-1       ] 
     739                     
     740                if cd_type == 'U' : 
     741                    ztab [..., -1, 0:-1     ] = psgn * ztab [..., -2, -2::-1       ]        
     742                    ztab [..., -1, -1       ] = psgn * ztab [..., -2, 0            ] # Bug ? 
     743                     
     744                if cd_type == 'V' : 
     745                    ztab [..., -1, 0:       ] = psgn * ztab [..., -3, -1::-1       ] 
     746                    ztab [..., -2, jpi//2:  ] = psgn * ztab [..., -2, jpi//2-1::-1 ] 
     747                     
     748                if cd_type == 'F' : 
     749                    ztab [..., -1, 0:-1     ] = psgn * ztab [..., -3, -2::-1       ] 
     750                    ztab [..., -1, -1       ] = psgn * ztab [..., -3, 0            ] 
     751                    ztab [..., -2, jpi//2:-1] = psgn * ztab [..., -2, jpi//2-2::-1 ] 
     752                     
     753            # 
     754            #> East-West boundary conditions 
     755            # ------------------------------ 
     756            if nperio in [1, 4, 6] : 
     757                # ... cyclic 
     758                ztab [...,  0] = ztab [..., -2] 
     759                ztab [..., -1] = ztab [...,  1] 
    642760 
    643761    if mmath == xr : 
    644         ztab = xr.DataArray ( ztab, dims=ptab.dims, coords=ptab.coords) 
     762        ztab = xr.DataArray ( ztab, dims=ptab.dims, coords=ptab.coords ) 
    645763        ztab.attrs = ptab.attrs 
    646764         
    647765    return ztab 
    648766 
    649 #@numba.jit(forceobj=True) 
    650767def lbc_mask (ptab, nperio=None, cd_type='T', sval=np.nan) : 
    651768    # 
    … …  
    659776    ''' 
    660777    jpj, jpi, nperio = lbc_init (ptab, nperio) 
     778    ix, ax = __findAxis__ (ptab, 'x') 
     779    jy, ay = __findAxis__ (ptab, 'y') 
    661780    ztab = ptab.copy () 
    662781 
    663     # 
    664     #> East-West boundary conditions 
    665     # ------------------------------ 
    666     if nperio in [1, 4, 6] : 
     782    if ax :  
     783        # 
     784        #> East-West boundary conditions 
     785        # ------------------------------ 
     786        if nperio in [1, 4, 6] : 
    667787        # ... cyclic 
    668         ztab [..., :,  0] = sval 
    669         ztab [..., :, -1] = sval 
    670  
    671     # 
    672     #> South (in which nperio cases ?) 
    673     # -------------------------------- 
    674     if nperio in [1, 3, 4, 5, 6] : 
    675         ztab [..., 0, :] = sval 
    676          
    677     # 
    678     #> North-South boundary conditions 
    679     # -------------------------------- 
    680     if nperio in [3, 4] :  # North fold T-point pivot 
    681         if cd_type in [ 'T', 'W' ] : # T-, W-point 
    682             ztab [..., -1,  :         ] = sval 
    683             ztab [..., -2, :jpi//2  ] = sval 
     788            ztab [...,  0] = sval 
     789            ztab [..., -1] = sval 
     790 
     791        if ay :  
     792            # 
     793            #> South (in which nperio cases ?) 
     794            # -------------------------------- 
     795            if nperio in [1, 3, 4, 5, 6] : 
     796                ztab [..., 0, :] = sval 
     797         
     798            # 
     799            #> North-South boundary conditions 
     800            # -------------------------------- 
     801            if nperio in [3, 4] :  # North fold T-point pivot 
     802                if cd_type in [ 'T', 'W' ] : # T-, W-point 
     803                    ztab [..., -1,  :         ] = sval 
     804                    ztab [..., -2, :jpi//2  ] = sval 
    684805                 
    685         if cd_type == 'U' : 
    686             ztab [..., -1,  :         ] = sval   
    687             ztab [..., -2, jpi//2+1:  ] = sval 
     806                if cd_type == 'U' : 
     807                    ztab [..., -1,  :         ] = sval   
     808                    ztab [..., -2, jpi//2+1:  ] = sval 
    688809                 
    689         if cd_type == 'V' : 
    690             ztab [..., -2, :       ] = sval 
    691             ztab [..., -1, :       ] = sval    
    692                  
    693         if cd_type == 'F' : 
    694             ztab [..., -2, :       ] = sval 
    695             ztab [..., -1, :       ] = sval 
    696  
    697     if nperio in [4.2] :  # North fold T-point pivot 
    698         if cd_type in [ 'T', 'W' ] : # T-, W-point 
    699             ztab [..., -1, jpi//2  :  ] = sval 
    700                  
    701         if cd_type == 'U' : 
    702             ztab [..., -1, jpi//2-1:-1] = sval 
    703                  
    704         if cd_type == 'V' :  
    705             ztab [..., -1, 1:       ] = sval 
    706                  
    707         if cd_type == 'F' : 
    708             ztab [..., -1, 0:-1     ] = sval 
    709      
    710     if nperio in [5, 6] :            #  North fold F-point pivot 
    711         if cd_type in ['T', 'W']  : 
    712             ztab [..., -1, 0:       ] = sval 
    713                  
    714         if cd_type == 'U' : 
    715             ztab [..., -1, 0:-1     ] = sval        
    716             ztab [..., -1, -1       ] = sval 
    717               
    718         if cd_type == 'V' : 
    719             ztab [..., -1, 0:       ] = sval 
    720             ztab [..., -2, jpi//2:  ] = sval 
    721                               
    722         if cd_type == 'F' : 
    723             ztab [..., -1, 0:-1       ] = sval 
    724             ztab [..., -1, -1         ] = sval 
    725             ztab [..., -2, jpi//2+1:-1] = sval 
     810                if cd_type == 'V' : 
     811                    ztab [..., -2, :       ] = sval 
     812                    ztab [..., -1, :       ] = sval    
     813 
     814                if cd_type == 'F' : 
     815                    ztab [..., -2, :       ] = sval 
     816                    ztab [..., -1, :       ] = sval 
     817 
     818            if nperio in [4.2] :  # North fold T-point pivot 
     819                if cd_type in [ 'T', 'W' ] : # T-, W-point 
     820                    ztab [..., -1, jpi//2  :  ] = sval 
     821 
     822                if cd_type == 'U' : 
     823                    ztab [..., -1, jpi//2-1:-1] = sval 
     824 
     825                if cd_type == 'V' :  
     826                    ztab [..., -1, 1:       ] = sval 
     827 
     828                if cd_type == 'F' : 
     829                    ztab [..., -1, 0:-1     ] = sval 
     830 
     831            if nperio in [5, 6] :            #  North fold F-point pivot 
     832                if cd_type in ['T', 'W']  : 
     833                    ztab [..., -1, 0:       ] = sval 
     834 
     835                if cd_type == 'U' : 
     836                    ztab [..., -1, 0:-1     ] = sval        
     837                    ztab [..., -1, -1       ] = sval 
     838 
     839                if cd_type == 'V' : 
     840                    ztab [..., -1, 0:       ] = sval 
     841                    ztab [..., -2, jpi//2:  ] = sval 
     842 
     843                if cd_type == 'F' : 
     844                    ztab [..., -1, 0:-1       ] = sval 
     845                    ztab [..., -1, -1         ] = sval 
     846                    ztab [..., -2, jpi//2+1:-1] = sval 
    726847 
    727848    return ztab 
    728849 
    729 #@numba.jit(forceobj=True) 
    730850def lbc_plot (ptab, nperio=None, cd_type='T', psgn=1.0, sval=np.nan) : 
    731851    ''' 
    … …  
    738858    See NEMO documentation for further details 
    739859    ''' 
    740  
    741860    jpj, jpi, nperio = lbc_init (ptab, nperio) 
     861    ix, ax = __findAxis__ (ptab, 'x') 
     862    jy, ay = __findAxis__ (ptab, 'y') 
    742863    psgn   = ptab.dtype.type (psgn) 
    743864    ztab   = ptab.copy () 
    744     # 
    745     #> East-West boundary conditions 
    746     # ------------------------------ 
    747     if nperio in [1, 4, 6] : 
    748         # ... cyclic 
    749         ztab [..., :,  0] = ztab [..., :, -2] 
    750         ztab [..., :, -1] = ztab [..., :,  1] 
    751  
    752     #> Masks south 
    753     # ------------ 
    754     if nperio in [4, 6] : ztab [..., 0, : ] = sval 
    755          
    756     # 
    757     #> North-South boundary conditions 
    758     # -------------------------------- 
    759     if nperio in [3, 4] :  # North fold T-point pivot 
    760         if cd_type in [ 'T', 'W' ] : # T-, W-point 
    761             ztab [..., -1,  :      ] = sval 
    762             #ztab [..., -2, jpi//2: ] = sval 
    763             ztab [..., -2, :jpi//2 ] = sval # Give better plots than above 
    764         if cd_type == 'U' : 
    765             ztab [..., -1, : ] = sval 
    766  
    767         if cd_type == 'V' :  
    768             ztab [..., -2, : ] = sval 
    769             ztab [..., -1, : ] = sval 
    770              
    771         if cd_type == 'F' : 
    772             ztab [..., -2, : ] = sval 
    773             ztab [..., -1, : ] = sval 
    774  
    775     if nperio in [4.2] :  # North fold T-point pivot 
    776         if cd_type in [ 'T', 'W' ] : # T-, W-point 
    777             ztab [..., -1, jpi//2:  ] = sval 
    778                  
    779         if cd_type == 'U' : 
    780             ztab [..., -1, jpi//2-1:-1] = sval 
    781                  
    782         if cd_type == 'V' :  
    783             ztab [..., -1, 1:       ] = sval 
    784                  
    785         if cd_type == 'F' : 
    786             ztab [..., -1, 0:-1     ] = sval 
    787        
    788     if nperio in [5, 6] :            #  North fold F-point pivot   
    789         if cd_type in ['T', 'W']  : 
    790             ztab [..., -1, : ] = sval 
    791                  
    792         if cd_type == 'U' : 
    793             ztab [..., -1, : ] = sval       
    794               
    795         if cd_type == 'V' : 
    796             ztab [..., -1, :        ] = sval 
    797             ztab [..., -2, jpi//2:  ] = sval 
    798                               
    799         if cd_type == 'F' : 
    800             ztab [..., -1, :          ] = sval 
    801             ztab [..., -2, jpi//2+1:-1] = sval 
     865 
     866    if ax :  
     867        # 
     868        #> East-West boundary conditions 
     869        # ------------------------------ 
     870        if nperio in [1, 4, 6] : 
     871            # ... cyclic 
     872            ztab [..., :,  0] = ztab [..., :, -2] 
     873            ztab [..., :, -1] = ztab [..., :,  1] 
     874 
     875        if ay :  
     876            #> Masks south 
     877            # ------------ 
     878            if nperio in [4, 6] : ztab [..., 0, : ] = sval 
     879 
     880            # 
     881            #> North-South boundary conditions 
     882            # -------------------------------- 
     883            if nperio in [3, 4] :  # North fold T-point pivot 
     884                if cd_type in [ 'T', 'W' ] : # T-, W-point 
     885                    ztab [..., -1,  :      ] = sval 
     886                    #ztab [..., -2, jpi//2: ] = sval 
     887                    ztab [..., -2, :jpi//2 ] = sval # Give better plots than above 
     888                if cd_type == 'U' : 
     889                    ztab [..., -1, : ] = sval 
     890 
     891                if cd_type == 'V' :  
     892                    ztab [..., -2, : ] = sval 
     893                    ztab [..., -1, : ] = sval 
     894 
     895                if cd_type == 'F' : 
     896                    ztab [..., -2, : ] = sval 
     897                    ztab [..., -1, : ] = sval 
     898 
     899            if nperio in [4.2] :  # North fold T-point pivot 
     900                if cd_type in [ 'T', 'W' ] : # T-, W-point 
     901                    ztab [..., -1, jpi//2:  ] = sval 
     902 
     903                if cd_type == 'U' : 
     904                    ztab [..., -1, jpi//2-1:-1] = sval 
     905 
     906                if cd_type == 'V' :  
     907                    ztab [..., -1, 1:       ] = sval 
     908 
     909                if cd_type == 'F' : 
     910                    ztab [..., -1, 0:-1     ] = sval 
     911 
     912            if nperio in [5, 6] :            #  North fold F-point pivot   
     913                if cd_type in ['T', 'W']  : 
     914                    ztab [..., -1, : ] = sval 
     915 
     916                if cd_type == 'U' : 
     917                    ztab [..., -1, : ] = sval       
     918 
     919                if cd_type == 'V' : 
     920                    ztab [..., -1, :        ] = sval 
     921                    ztab [..., -2, jpi//2:  ] = sval 
     922 
     923                if cd_type == 'F' : 
     924                    ztab [..., -1, :          ] = sval 
     925                    ztab [..., -2, jpi//2+1:-1] = sval 
    802926 
    803927    return ztab 
    804928 
    805 #@numba.jit(forceobj=True) 
    806929def lbc_add (ptab, nperio=None, cd_type=None, psgn=1, sval=None) : 
    807930    ''' 
    808     Handle NEMO domain changes between NEMO 4.0 to NEMO 4.2 
     931    Handles NEMO domain changes between NEMO 4.0 to NEMO 4.2 
    809932      Peridodicity halo has been removed 
    810933    This routine adds the halos if needed 
    … …  
    818941    mmath = __mmath__ (ptab)  
    819942    jpj, jpi, nperio = lbc_init (ptab, nperio) 
     943    lshape = get_shape (ptab) 
     944    ix, ax = __findAxis__ (ptab, 'x') 
     945    jy, ay = __findAxis__ (ptab, 'y') 
    820946 
    821947    t_shape = np.array (ptab.shape) 
    … …  
    823949    if nperio == 4.2 or nperio == 6.2 : 
    824950       
    825         ext_shape = t_shape 
    826         ext_shape[-1] = ext_shape[-1] + 2 
    827         ext_shape[-2] = ext_shape[-2] + 1 
     951        ext_shape = t_shape.copy() 
     952        if 'X' in lshape : ext_shape[ix] = ext_shape[ix] + 2 
     953        if 'Y' in lshape : ext_shape[jy] = ext_shape[jy] + 1 
    828954 
    829955        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 () 
     956            ptab_ext = xr.DataArray (np.zeros (ext_shape), dims=ptab.dims) 
     957            if 'X' in lshape and 'Y' in lshape : 
     958                ptab_ext.values[..., :-1, 1:-1] = ptab.values.copy () 
     959            else : 
     960                if 'X' in lshape     : ptab_ext.values[...,      1:-1] = ptab.values.copy () 
     961                if 'Y' in lshape     : ptab_ext.values[..., :-1      ] = ptab.values.copy () 
    832962        else           : 
    833963            ptab_ext =               np.zeros (ext_shape) 
    834             ptab_ext[..., :-1, 1:-1] = ptab.copy () 
    835              
    836         #if sval != None :  ptab_ext[..., 0, :] = sval 
    837          
     964            if 'X' in lshape and 'Y' in lshape : ptab_ext       [..., :-1, 1:-1] = ptab.copy () 
     965            else : 
     966                if 'X' in lshape     : ptab_ext       [...,      1:-1] = ptab.copy () 
     967                if 'Y' in lshape     : ptab_ext       [..., :-1      ] = ptab.copy ()             
     968 
    838969        if nperio == 4.2 : ptab_ext = lbc (ptab_ext, nperio=4, cd_type=cd_type, psgn=psgn) 
    839970        if nperio == 6.2 : ptab_ext = lbc (ptab_ext, nperio=6, cd_type=cd_type, psgn=psgn) 
    840               
     971         
    841972        if mmath == xr : 
    842973            ptab_ext.attrs = ptab.attrs 
     974            kz, az = __findAxis__ (ptab, 'z') 
     975            it, at = __findAxis__ (ptab, 't') 
     976            if az : ptab_ext = ptab_ext.assign_coords ( {az:ptab.coords[az]} ) 
     977            if at : ptab_ext = ptab_ext.assign_coords ( {at:ptab.coords[at]} ) 
    843978 
    844979    else : ptab_ext = lbc (ptab, nperio=nperio, cd_type=cd_type, psgn=psgn) 
    … …  
    848983def lbc_del (ptab, nperio=None, cd_type='T', psgn=1) : 
    849984    ''' 
    850     Handle NEMO domain changes between NEMO 4.0 to NEMO 4.2 
     985    Handles NEMO domain changes between NEMO 4.0 to NEMO 4.2 
    851986      Periodicity halo has been removed 
    852987    This routine removes the halos if needed 
    … …  
    858993    See NEMO documentation for further details 
    859994    ''' 
    860  
    861995    jpj, jpi, nperio = lbc_init (ptab, nperio) 
     996    lshape = get_shape (ptab) 
     997    ix, ax = __findAxis__ (ptab, 'x') 
     998    jy, ay = __findAxis__ (ptab, 'y') 
    862999 
    8631000    if nperio == 4.2 or nperio == 6.2 : 
    864         return lbc (ptab[..., :-1, 1:-1], nperio=nperio, cd_type=cd_type, psgn=psgn) 
     1001        if ax or ay :  
     1002            if ax and ay :  
     1003                return lbc (ptab[..., :-1, 1:-1], nperio=nperio, cd_type=cd_type, psgn=psgn) 
     1004            else :  
     1005                if ax : 
     1006                    return lbc (ptab[...,      1:-1], nperio=nperio, cd_type=cd_type, psgn=psgn) 
     1007                if ay : 
     1008                    return lbc (ptab[..., -1], nperio=nperio, cd_type=cd_type, psgn=psgn) 
     1009        else : 
     1010            return ptab 
    8651011    else : 
    8661012        return ptab 
    8671013 
    868 #@numba.jit(forceobj=True) 
    8691014def lbc_index (jj, ii, jpj, jpi, nperio=None, cd_type='T') : 
    8701015    ''' 
    … …  
    9511096    return jy, ix 
    9521097     
    953 #def geo2en (pxx, pyy, pzz, glam, gphi) :  
     1098def findJI (lat_data, lon_data, lat_grid, lon_grid, mask=1.0, verbose=False, out=None) : 
     1099    ''' 
     1100    Description: seeks J,I indices of the grid point which is the closest of a given point  
     1101    Usage: go FindJI  <data latitude> <data longitude> <grid latitudes> <grid longitudes> [mask] 
     1102    <longitude fields> <latitude field> are 2D fields on J/I (Y/X) dimensions 
     1103    mask : if given, seek only non masked grid points (i.e with mask=1) 
     1104     
     1105    Example : findIJ (40, -20, nav_lat, nav_lon, mask=1.0) 
     1106 
     1107    Note : all longitudes and latitudes in degrees 
     1108         
     1109    Note : may work with 1D lon/lat (?) 
     1110    ''' 
     1111    # Get grid dimensions 
     1112    if len (lon_grid.shape) == 2 : (jpj, jpi) = lon_grid.shape 
     1113    else                         : jpj = len(lat_grid) ; jpi=len(lon_grid) 
     1114 
     1115    mmath = __mmath__ (lat_grid) 
     1116         
     1117    # Compute distance from the point to all grid points (in radian) 
     1118    arg      = np.sin (rad*lat_data) * np.sin (rad*lat_grid) \ 
     1119             + np.cos (rad*lat_data) * np.cos (rad*lat_grid) * np.cos(rad*(lon_data-lon_grid)) 
     1120    distance = np.arccos (arg) + 4.0*rpi*(1.0-mask) # Send masked points to 'infinite'  
     1121 
     1122    # Truncates to alleviate some precision problem with some grids 
     1123    prec = int (1E7) 
     1124    distance = (distance*prec).astype(int) / prec 
     1125 
     1126    # Compute minimum of distance, and index of minimum 
     1127    # 
     1128    distance_min = distance.min    () 
     1129    jimin        = int (distance.argmin ()) 
     1130     
     1131    # Compute 2D indices  
     1132    jmin = jimin // jpi ; imin = jimin - jmin*jpi 
     1133     
     1134    # Result 
     1135    if verbose : 
     1136        # Compute distance achieved 
     1137        mindist = distance [jmin, imin] 
     1138         
     1139        # Compute azimuth 
     1140        dlon = lon_data-lon_grid[jmin,imin] 
     1141        arg  = np.sin (rad*dlon) /  (np.cos(rad*lat_data)*np.tan(rad*lat_grid[jmin,imin]) - np.sin(rad*lat_data)*np.cos(rad*dlon)) 
     1142        azimuth = dar*np.arctan (arg) 
     1143        print ( f'I={imin:d} J={jmin:d} - Data:{lat_data:5.1f}°N {lon_data:5.1f}°E - Grid:{lat_grid[jmin,imin]:4.1f}°N '   \ 
     1144            +   f'{lon_grid[jmin,imin]:4.1f}°E - Dist: {ra*distance[jmin,imin]:6.1f}km {dar*distance[jmin,imin]:5.2f}° ' \ 
     1145            +   f'- Azimuth: {rad*azimuth:3.2f}rad - {azimuth:5.1f}°' ) 
     1146 
     1147    if   out=='dict'                               : return {'x':imin, 'y':jmin} 
     1148    elif out=='array' or out=='numpy'  or out=='np': return np.array ( [jmin, imin] ) 
     1149    elif out=='xarray' or out=='xr'                : return xr.DataArray ( [jmin, imin] ) 
     1150    elif out=='list'                               : return [jmin, imin] 
     1151    elif out=='tuple'                              : return jmin, imin 
     1152    else                                           : return jmin, imin 
     1153 
     1154def geo2en (pxx, pyy, pzz, glam, gphi) :  
    9541155    ''' 
    9551156    Change vector from geocentric to east/north 
    … …  
    9701171    return pte, ptn 
    9711172 
    972 #def en2geo (pte, ptn, glam, gphi) : 
     1173def en2geo (pte, ptn, glam, gphi) : 
    9731174    ''' 
    9741175    Change vector from east/north to geocentric 
    … …  
    9901191    return pxx, pyy, pzz 
    9911192 
    992 #def findJI (lat_data, lon_data, lat_grid, lon_grid, mask=1.0, verbose=False) : 
    993     ''' 
    994     Description: seeks J,I indices of the grid point which is the closest of a given point  
    995     Usage: go FindJI  <data latitude> <data longitude> <grid latitudes> <grid longitudes> [mask] 
    996     <longitude fields> <latitude field> are 2D fields on J/I (Y/X) dimensions 
    997     mask : if given, seek only non masked grid points (i.e with mask=1) 
    998      
    999     Example : findIJ (40, -20, nav_lat, nav_lon, mask=1.0) 
    1000  
    1001     Note : all longitudes and latitudes in degrees 
    1002          
    1003     Note : may work with 1D lon/lat (?) 
    1004     ''' 
    1005     # Get grid dimensions 
    1006     if len (lon_grid.shape) == 2 : (jpj, jpi) = lon_grid.shape 
    1007     else                         : jpj = len(lat_grid) ; jpi=len(lon_grid) 
    1008  
    1009     mmath = __mmath__ (lat_grid) 
    1010          
    1011     # Compute distance from the point to all grid points (in radian) 
    1012     arg      = np.sin (rad*lat_data) * np.sin (rad*lat_grid) \ 
    1013              + np.cos (rad*lat_data) * np.cos (rad*lat_grid) * np.cos(rad*(lon_data-lon_grid)) 
    1014     distance = np.arccos (arg) + 4.0*rpi*(1.0-mask) # Send masked points to 'infinite'  
    1015  
    1016     # Truncates to alleviate some precision problem with some grids 
    1017     prec = int (1E7) 
    1018     distance = (distance*prec).astype(int) / prec 
    1019  
    1020     # Compute minimum of distance, and index of minimum 
    1021     # 
    1022     distance_min = distance.min    () 
    1023     jimin        = int (distance.argmin ()) 
    1024      
    1025     # Compute 2D indices  
    1026     jmin = jimin // jpi ; imin = jimin - jmin*jpi 
    1027  
    1028     # Compute distance achieved 
    1029     mindist = distance[jmin, imin] 
    1030      
    1031     # Compute azimuth 
    1032     dlon = lon_data-lon_grid[jmin,imin] 
    1033     arg  = np.sin (rad*dlon) /  (np.cos(rad*lat_data)*np.tan(rad*lat_grid[jmin,imin]) - np.sin(rad*lat_data)*np.cos(rad*dlon)) 
    1034     azimuth = dar*np.arctan (arg) 
    1035      
    1036     # Result 
    1037     if verbose :  
    1038         print ('I={:d} J={:d} - Data:{:5.1f}°N {:5.1f}°E - Grid:{:4.1f}°N {:4.1f}°E - Dist: {:6.1f}km {:5.2f}° - Azimuth: {:3.2f}rad - {:5.1f}°' 
    1039             .format (imin, jmin, lat_data, lon_data, lat_grid[jmin,imin], lon_grid[jmin,imin], ra*distance[jmin,imin], dar*distance[jmin,imin], rad*azimuth, azimuth)) 
    1040  
    1041     return jmin, imin 
    1042  
    1043 #def clo_lon (lon, lon0) : 
     1193 
     1194def clo_lon (lon, lon0=0., rad=False, deg=True) : 
    10441195    '''Choose closest to lon0 longitude, adding or substacting 360° if needed''' 
    10451196    mmath = __mmath__ (lon, np) 
    1046          
     1197    if rad : lon_range = 2.*np.pi 
     1198    if deg : lon_range = 360. 
    10471199    clo_lon = 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) 
     1200    clo_lon = mmath.where (clo_lon > lon0 + lon_range*0.5, clo_lon-lon_range, clo_lon) 
     1201    clo_lon = mmath.where (clo_lon < lon0 - lon_range*0.5, clo_lon+lon_range, clo_lon) 
     1202    clo_lon = mmath.where (clo_lon > lon0 + lon_range*0.5, clo_lon-lon_range, clo_lon) 
     1203    clo_lon = mmath.where (clo_lon < lon0 - lon_range*0.5, clo_lon+lon_range, clo_lon) 
    10521204    if clo_lon.shape == () : clo_lon = clo_lon.item () 
     1205    if mmath == xr : 
     1206        try :  
     1207            for attr in lon.attrs : clo_lon.attrs[attr] = lon.attrs[attr] 
     1208        except : 
     1209            pass 
    10531210    return clo_lon 
    10541211 
    1055 #def angle_full (glamt, gphit, glamu, gphiu, glamv, gphiv, glamf, gphif, nperio=None) : 
     1212 
     1213def index2depth (pk, gdept_0) : 
     1214    ''' 
     1215    From index (real, continuous), get depth 
     1216    ''' 
     1217    jpk = gdept_0.shape[0] 
     1218    kk = xr.DataArray(pk) 
     1219    k  = np.maximum (0, np.minimum (jpk-1, kk    )) 
     1220    k0 = np.floor (k).astype (int) 
     1221    k1 = np.maximum (0, np.minimum (jpk-1,  k0+1)) 
     1222    zz = k - k0 
     1223    gz = (1.0-zz)*gdept_0[k0]+ zz*gdept_0[k1] 
     1224    return gz.values 
     1225 
     1226def depth2index (pz, gdept_0) : 
     1227    ''' 
     1228    From depth, get index (real, continuous) 
     1229    ''' 
     1230    jpk  = gdept_0.shape[0] 
     1231    if type (pz) == xr.core.dataarray.DataArray : 
     1232        zz   = xr.DataArray (pz.values, dims=('zz',)) 
     1233    elif type (pz) == np.ndarray : 
     1234        zz   = xr.DataArray (pz.ravel(), dims=('zz',)) 
     1235    else : 
     1236        zz   = xr.DataArray (np.array([pz]).ravel(), dims=('zz',)) 
     1237    zz   = np.minimum (gdept_0[-1], np.maximum (0, zz)) 
     1238     
     1239    idk1 = np.minimum ( (gdept_0-zz), 0.).argmax (axis=0).astype(int) 
     1240    idk1 = np.maximum (0, np.minimum (jpk-1,  idk1  )) 
     1241    idk2 = np.maximum (0, np.minimum (jpk-1,  idk1-1)) 
     1242     
     1243    ff = (zz - gdept_0[idk2])/(gdept_0[idk1]-gdept_0[idk2]) 
     1244    idk = ff*idk1 + (1.0-ff)*idk2 
     1245    idk = xr.where ( np.isnan(idk), idk1, idk) 
     1246    return idk.values 
     1247 
     1248def index2depth_panels (pk, gdept_0, depth0, fact) : 
     1249    ''' 
     1250    From  index (real, continuous), get depth, with bottom part compressed 
     1251    ''' 
     1252    jpk = gdept_0.shape[0] 
     1253    kk = xr.DataArray (pk) 
     1254    k  = np.maximum (0, np.minimum (jpk-1, kk    )) 
     1255    k0 = np.floor (k).astype (int) 
     1256    k1 = np.maximum (0, np.minimum (jpk-1,  k0+1)) 
     1257    zz = k - k0 
     1258    gz = (1.0-zz)*gdept_0[k0]+ zz*gdept_0[k1] 
     1259    gz = xr.where ( gz<depth0, gz, depth0 + (gz-depth0)*fact) 
     1260    return gz.values 
     1261 
     1262def depth2index_panels (pz, gdept_0, depth0, fact) : 
     1263    ''' 
     1264    From  index (real, continuous), get depth, with bottom part compressed 
     1265    ''' 
     1266    jpk = gdept_0.shape[0] 
     1267    if type (pz) == xr.core.dataarray.DataArray : 
     1268        zz   = xr.DataArray (pz.values , dims=('zz',)) 
     1269    elif type (pz) == np.ndarray : 
     1270        zz   = xr.DataArray (pz.ravel(), dims=('zz',)) 
     1271    else :  
     1272        zz   = xr.DataArray (np.array([pz]).ravel(), dims=('zz',)) 
     1273    zz         = np.minimum (gdept_0[-1], np.maximum (0, zz)) 
     1274    gdept_comp = xr.where ( gdept_0>depth0, (gdept_0-depth0)*fact+depth0, gdept_0) 
     1275    zz_comp    = xr.where ( zz     >depth0, (zz     -depth0)*fact+depth0, zz     ) 
     1276    zz_comp    = np.minimum (gdept_comp[-1], np.maximum (0, zz_comp)) 
     1277 
     1278    idk1 = np.minimum ( (gdept_0-zz_comp), 0.).argmax (axis=0).astype(int) 
     1279    idk1 = np.maximum (0, np.minimum (jpk-1,  idk1  )) 
     1280    idk2 = np.maximum (0, np.minimum (jpk-1,  idk1-1)) 
     1281      
     1282    ff = (zz_comp - gdept_0[idk2])/(gdept_0[idk1]-gdept_0[idk2]) 
     1283    idk = ff*idk1 + (1.0-ff)*idk2 
     1284    idk = xr.where ( np.isnan(idk), idk1, idk) 
     1285    return idk.values 
     1286 
     1287def depth2comp (pz, depth0, fact ) : 
     1288    ''' 
     1289    Form depth, get compressed depth, with bottom part compressed 
     1290    ''' 
     1291    #print ('start depth2comp') 
     1292    if type (pz) == xr.core.dataarray.DataArray : 
     1293        zz   = pz.values 
     1294    elif type(pz) == list : 
     1295        zz = np.array (pz) 
     1296    else :  
     1297        zz   = pz 
     1298    gz = np.where ( zz>depth0, (zz-depth0)*fact+depth0, zz) 
     1299    #print ( f'depth2comp : {gz=}' ) 
     1300    if type (pz) in [int, float] : return gz.item() 
     1301    else : return gz 
     1302    #print ('end depth2comp') 
     1303 
     1304def comp2depth (pz, depth0, fact ) : 
     1305    ''' 
     1306    Form compressed depth, get depth, with bottom part compressed 
     1307    ''' 
     1308    if type (pz) == xr.core.dataarray.DataArray : 
     1309        zz   = pz.values 
     1310    elif type(pz) == list : 
     1311        zz = np.array (pz) 
     1312    else :  
     1313        zz   = pz 
     1314    gz = np.where ( zz>depth0, (zz-depth0)/fact+depth0, zz) 
     1315    if type (pz) in [int, float] : return gz.item() 
     1316    else : return gz 
     1317 
     1318def index2lon (pi, lon1D) : 
     1319    ''' 
     1320    From index (real, continuous), get longitude 
     1321    ''' 
     1322    jpi = lon1D.shape[0] 
     1323    ii = xr.DataArray (pi) 
     1324    i =  np.maximum (0, np.minimum (jpi-1, ii    )) 
     1325    i0 = np.floor (i).astype (int) 
     1326    i1 = np.maximum (0, np.minimum (jpi-1,  i0+1)) 
     1327    xx = i - i0 
     1328    gx = (1.0-xx)*lon1D[i0]+ xx*lon1D[i1] 
     1329    return gx.values 
     1330 
     1331def lon2index (px, lon1D) : 
     1332    ''' 
     1333    From longitude, get index (real, continuous) 
     1334    ''' 
     1335    jpi  = lon1D.shape[0] 
     1336    if type (px) == xr.core.dataarray.DataArray : 
     1337        xx   = xr.DataArray (px.values , dims=('xx',)) 
     1338    elif type (px) == np.ndarray : 
     1339        xx   = xr.DataArray (px.ravel(), dims=('xx',)) 
     1340    else :  
     1341        xx   = xr.DataArray (np.array([px]).ravel(), dims=('xx',)) 
     1342    xx   = xr.where ( xx>lon1D.max(), xx-360.0, xx) 
     1343    xx   = xr.where ( xx<lon1D.min(), xx+360.0, xx) 
     1344    xx   = np.minimum (lon1D.max(), np.maximum(xx, lon1D.min() )) 
     1345    idi1 = np.minimum ( (lon1D-xx), 0.).argmax (axis=0).astype(int) 
     1346    idi1 = np.maximum (0, np.minimum (jpi-1,  idi1  )) 
     1347    idi2 = np.maximum (0, np.minimum (jpi-1,  idi1-1)) 
     1348     
     1349    ff = (xx - lon1D[idi2])/(lon1D[idi1]-lon1D[idi2]) 
     1350    idi = ff*idi1 + (1.0-ff)*idi2 
     1351    idi = xr.where ( np.isnan(idi), idi1, idi) 
     1352    return idi.values 
     1353 
     1354def index2lat (pj, lat1D) : 
     1355    ''' 
     1356    From index (real, continuous), get latitude 
     1357    ''' 
     1358    jpj = lat1D.shape[0] 
     1359    jj  = xr.DataArray (pj) 
     1360    j   = np.maximum (0, np.minimum (jpj-1, jj    )) 
     1361    j0  = np.floor (j).astype (int) 
     1362    j1  = np.maximum (0, np.minimum (jpj-1,  j0+1)) 
     1363    yy  = j - j0 
     1364    gy  = (1.0-yy)*lat1D[j0]+ yy*lat1D[j1] 
     1365    return gy.values 
     1366 
     1367def lat2index (py, lat1D) : 
     1368    ''' 
     1369    From latitude, get index (real, continuous) 
     1370    ''' 
     1371    jpj = lat1D.shape[0] 
     1372    if type (py) == xr.core.dataarray.DataArray : 
     1373        yy   = xr.DataArray (py.values , dims=('yy',)) 
     1374    elif type (py) == np.ndarray : 
     1375        yy   = xr.DataArray (py.ravel(), dims=('yy',)) 
     1376    else :  
     1377        yy   = xr.DataArray (np.array([py]).ravel(), dims=('yy',)) 
     1378    yy   = np.minimum (lat1D.max(), np.minimum(yy, lat1D.max() )) 
     1379    idj1 = np.minimum ( (lat1D-yy), 0.).argmax (axis=0).astype(int) 
     1380    idj1 = np.maximum (0, np.minimum (jpj-1,  idj1  )) 
     1381    idj2 = np.maximum (0, np.minimum (jpj-1,  idj1-1)) 
     1382     
     1383    ff = (yy - lat1D[idj2])/(lat1D[idj1]-lat1D[idj2]) 
     1384    idj = ff*idj1 + (1.0-ff)*idj2 
     1385    idj = xr.where ( np.isnan(idj), idj1, idj) 
     1386    return idj.values 
     1387 
     1388def angle_full (glamt, gphit, glamu, gphiu, glamv, gphiv, glamf, gphif, nperio=None) : 
    10561389    '''Compute sinus and cosinus of model line direction with respect to east''' 
    10571390    mmath = __mmath__ (glamt) 
    … …  
    12341567    v_n = + u_i * gsin + v_j * gcos 
    12351568     
    1236     u_e = lbc (u_e, nperio=nperio, cd_type=cd_type, psgn= 1.0) 
    1237     v_n = lbc (v_n, nperio=nperio, cd_type=cd_type, psgn= 1.0) 
     1569    u_e = lbc (u_e, nperio=nperio, cd_type=cd_type, psgn=1.0) 
     1570    v_n = lbc (v_n, nperio=nperio, cd_type=cd_type, psgn=1.0) 
    12381571     
    12391572    return u_e, v_n 
    12401573 
    1241 def rot_uv2en ( uo, vo, gsint, gcost, nperio, zdim='deptht' ) : 
     1574def rot_uv2en ( uo, vo, gsint, gcost, nperio, zdim=None ) : 
    12421575    ''' 
    12431576    ** Purpose :   Rotate the Repere: Change vector componantes from 
    … …  
    12471580    ''' 
    12481581    mmath = __mmath__ (uo) 
    1249  
     1582     
    12501583    ut = U2T (uo, nperio=nperio, psgn=-1.0, zdim=zdim) 
    12511584    vt = V2T (vo, nperio=nperio, psgn=-1.0, zdim=zdim) 
    … …  
    12591592    return u_e, v_n 
    12601593 
    1261 def rot_uv2enF ( uo, vo, gsinf, gcosf, nperio, zdim='deptht' ) : 
     1594def rot_uv2enF ( uo, vo, gsinf, gcosf, nperio, zdim=None ) : 
    12621595    ''' 
    12631596    ** Purpose : Rotate the Repere: Change vector componantes from 
    … …  
    12791612    return u_e, v_n 
    12801613 
    1281 #@numba.jit(forceobj=True) 
    1282 def U2T (utab, nperio=None, psgn=-1.0, zdim='deptht', action='ave') : 
     1614def U2T (utab, nperio=None, psgn=-1.0, zdim=None, action='ave') : 
    12831615    '''Interpolate an array from U grid to T grid i-mean)''' 
    12841616    mmath = __mmath__ (utab) 
    … …  
    12871619    utab_0 = lbc_add (utab_0, nperio=nperio, cd_type='U', psgn=psgn) 
    12881620    ix, ax = __findAxis__ (utab_0, 'x') 
    1289     iz, az = __findAxis__ (utab_0, 'z') 
    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      
     1621    kz, az = __findAxis__ (utab_0, 'z') 
     1622 
     1623    if ax :  
     1624        if action == 'ave' : ttab = 0.5 *      (utab_0 + np.roll (utab_0, axis=ix, shift=1)) 
     1625        if action == 'min' : ttab = np.minimum (utab_0 , np.roll (utab_0, axis=ix, shift=1)) 
     1626        if action == 'max' : ttab = np.maximum (utab_0 , np.roll (utab_0, axis=ix, shift=1)) 
     1627        if action == 'mult': ttab =             utab_0 * np.roll (utab_0, axis=ix, shift=1) 
     1628        ttab = lbc_del (ttab  , nperio=nperio, cd_type='T', psgn=psgn) 
     1629    else :  
     1630        ttab = lbc_del (utab_0, nperio=nperio, cd_type='T', psgn=psgn) 
     1631         
    12961632    if mmath == xr : 
    1297         if ax != None : 
     1633        if ax : 
    12981634            ttab = ttab.assign_coords({ax:np.arange (ttab.shape[ix])+1.}) 
    1299         if zdim != None and iz != None  and az != 'olevel' :  
    1300             ttab = ttab.rename( {az:zdim})  
     1635        if zdim and az : 
     1636            if az != zdim : ttab = ttab.rename( {az:zdim})  
    13011637    return ttab 
    13021638 
    1303 #@numba.jit(forceobj=True) 
    1304 def V2T (vtab, nperio=None, psgn=-1.0, zdim='deptht', action='ave') : 
     1639def V2T (vtab, nperio=None, psgn=-1.0, zdim=None, action='ave') : 
    13051640    '''Interpolate an array from V grid to T grid (j-mean)''' 
    13061641    mmath = __mmath__ (vtab) 
    … …  
    13081643    vtab_0 = mmath.where ( np.isnan(vtab), 0., vtab) 
    13091644    vtab_0 = lbc_add (vtab_0, nperio=nperio, cd_type='V', psgn=psgn) 
    1310     iy, ay = __findAxis__ (vtab_0, 'y') 
    1311     iz, az = __findAxis__ (vtab_0, 'z') 
    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) 
     1645    jy, ay = __findAxis__ (vtab_0, 'y') 
     1646    kz, az = __findAxis__ (vtab_0, 'z') 
     1647    if ay :  
     1648        if action == 'ave'  : ttab = 0.5 *      (vtab_0 + np.roll (vtab_0, axis=jy, shift=1)) 
     1649        if action == 'min'  : ttab = np.minimum (vtab_0 , np.roll (vtab_0, axis=jy, shift=1)) 
     1650        if action == 'max'  : ttab = np.maximum (vtab_0 , np.roll (vtab_0, axis=jy, shift=1)) 
     1651        if action == 'mult' : ttab =             vtab_0 * np.roll (vtab_0, axis=jy, shift=1) 
     1652        ttab = lbc_del (ttab  , nperio=nperio, cd_type='T', psgn=psgn) 
     1653    else : 
     1654        ttab = lbc_del (vtab_0, nperio=nperio, cd_type='T', psgn=psgn) 
     1655         
    13171656    if mmath == xr : 
    1318         if ay !=None :  
    1319             ttab = ttab.assign_coords({ay:np.arange(ttab.shape[iy])+1.}) 
    1320         if zdim != None and iz != None  and az != 'olevel' : 
    1321             ttab = ttab.rename( {az:zdim})  
     1657        if ay : 
     1658            ttab = ttab.assign_coords({ay:np.arange(ttab.shape[jy])+1.}) 
     1659        if zdim and az : 
     1660            if az != zdim : ttab = ttab.rename( {az:zdim})  
    13221661    return ttab 
    13231662 
    1324 #@numba.jit(forceobj=True) 
    1325 def F2T (ftab, nperio=None, psgn=1.0, zdim='depthf', action='ave') : 
     1663def F2T (ftab, nperio=None, psgn=1.0, zdim=None, action='ave') : 
    13261664    '''Interpolate an array from F grid to T grid (i- and j- means)''' 
    13271665    mmath = __mmath__ (ftab) 
    13281666    ftab_0 = mmath.where ( np.isnan(ftab), 0., ftab) 
    13291667    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) 
     1668    ttab = V2T (F2V (ftab_0, nperio=nperio, psgn=psgn, zdim=zdim, action=action), nperio=nperio, psgn=psgn, zdim=zdim, action=action) 
    13311669    return lbc_del (ttab, nperio=nperio, cd_type='T', psgn=psgn) 
    13321670 
    1333 #@numba.jit(forceobj=True) 
    1334 def T2U (ttab, nperio=None, psgn=1.0, zdim='depthu', action='ave') : 
     1671def T2U (ttab, nperio=None, psgn=1.0, zdim=None, action='ave') : 
    13351672    '''Interpolate an array from T grid to U grid (i-mean)''' 
    13361673    mmath = __mmath__ (ttab) 
    … …  
    13381675    ttab_0 = lbc_add (ttab_0 , nperio=nperio, cd_type='T', psgn=psgn) 
    13391676    ix, ax = __findAxis__ (ttab_0, 'x') 
    1340     iz, az = __findAxis__ (ttab_0, 'z') 
    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  
     1677    kz, az = __findAxis__ (ttab_0, 'z') 
     1678    if ix :  
     1679        if action == 'ave'  : utab = 0.5 *      (ttab_0 + np.roll (ttab_0, axis=ix, shift=-1)) 
     1680        if action == 'min'  : utab = np.minimum (ttab_0 , np.roll (ttab_0, axis=ix, shift=-1)) 
     1681        if action == 'max'  : utab = np.maximum (ttab_0 , np.roll (ttab_0, axis=ix, shift=-1)) 
     1682        if action == 'mult' : utab =             ttab_0 * np.roll (ttab_0, axis=ix, shift=-1) 
     1683        utab = lbc_del (utab  , nperio=nperio, cd_type='U', psgn=psgn) 
     1684    else : 
     1685        utab = lbc_del (ttab_0, nperio=nperio, cd_type='U', psgn=psgn) 
     1686         
    13471687    if mmath == xr :     
    1348         if ax != None :  
     1688        if ax :  
    13491689            utab = ttab.assign_coords({ax:np.arange(utab.shape[ix])+1.}) 
    1350         if zdim != None  and iz != None  and az != 'olevel' : 
    1351             utab = utab.rename( {az:zdim})  
     1690        if zdim and az : 
     1691            if az != zdim : utab = utab.rename( {az:zdim})  
    13521692    return utab 
    13531693 
    1354 #@numba.jit(forceobj=True) 
    1355 def T2V (ttab, nperio=None, psgn=1.0, zdim='depthv', action='ave') : 
     1694def T2V (ttab, nperio=None, psgn=1.0, zdim=None, action='ave') : 
    13561695    '''Interpolate an array from T grid to V grid (j-mean)''' 
    13571696    mmath = __mmath__ (ttab) 
    13581697    ttab_0 = mmath.where ( np.isnan(ttab), 0., ttab) 
    13591698    ttab_0 = lbc_add (ttab_0 , nperio=nperio, cd_type='T', psgn=psgn) 
    1360     iy, ay = __findAxis__ (ttab_0, 'y') 
    1361     iz, az = __findAxis__ (ttab_0, 'z') 
    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) 
     1699    jy, ay = __findAxis__ (ttab_0, 'y') 
     1700    kz, az = __findAxis__ (ttab_0, 'z') 
     1701    if jy :  
     1702        if action == 'ave'  : vtab = 0.5 *      (ttab_0 + np.roll (ttab_0, axis=jy, shift=-1)) 
     1703        if action == 'min'  : vtab = np.minimum (ttab_0 , np.roll (ttab_0, axis=jy, shift=-1)) 
     1704        if action == 'max'  : vtab = np.maximum (ttab_0 , np.roll (ttab_0, axis=jy, shift=-1)) 
     1705        if action == 'mult' : vtab =             ttab_0 * np.roll (ttab_0, axis=jy, shift=-1) 
     1706        vtab = lbc_del (vtab  , nperio=nperio, cd_type='V', psgn=psgn) 
     1707    else : 
     1708        vtab = lbc_del (ttab_0, nperio=nperio, cd_type='V', psgn=psgn) 
     1709 
    13681710    if mmath == xr : 
    1369         if ay != None :  
    1370             vtab = vtab.assign_coords({ay:np.arange(vtab.shape[iy])+1.}) 
    1371         if zdim != None  and iz != None and az != 'olevel' : 
    1372             vtab = vtab.rename( {az:zdim})  
     1711        if ay :  
     1712            vtab = vtab.assign_coords({ay:np.arange(vtab.shape[jy])+1.}) 
     1713        if zdim and az : 
     1714            if az != zdim : vtab = vtab.rename( {az:zdim})  
    13731715    return vtab 
    13741716 
    1375 #@numba.jit(forceobj=True) 
    1376 def V2F (vtab, nperio=None, psgn=-1.0, zdim='depthf', action='ave') : 
     1717def V2F (vtab, nperio=None, psgn=-1.0, zdim=None, action='ave') : 
    13771718    '''Interpolate an array from V grid to F grid (i-mean)''' 
    13781719    mmath = __mmath__ (vtab) 
    … …  
    13801721    vtab_0 = lbc_add (vtab_0 , nperio=nperio, cd_type='V', psgn=psgn) 
    13811722    ix, ax = __findAxis__ (vtab_0, 'x') 
    1382     iz, az = __findAxis__ (vtab_0, 'z') 
    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      
     1723    kz, az = __findAxis__ (vtab_0, 'z') 
     1724    if ix :  
     1725        if action == 'ave'  : 0.5 *      (vtab_0 + np.roll (vtab_0, axis=ix, shift=-1)) 
     1726        if action == 'min'  : np.minimum (vtab_0 , np.roll (vtab_0, axis=ix, shift=-1)) 
     1727        if action == 'max'  : np.maximum (vtab_0 , np.roll (vtab_0, axis=ix, shift=-1)) 
     1728        if action == 'mult' :             vtab_0 * np.roll (vtab_0, axis=ix, shift=-1) 
     1729        ftab = lbc_del (ftab  , nperio=nperio, cd_type='F', psgn=psgn) 
     1730    else : 
     1731         ftab = lbc_del (vtab_0, nperio=nperio, cd_type='F', psgn=psgn) 
     1732    
    13891733    if mmath == xr : 
    1390         if ax != None :  
     1734        if ax :  
    13911735            ftab = ftab.assign_coords({ax:np.arange(ftab.shape[ix])+1.}) 
    1392         if zdim != None and iz != None and az != 'olevel' : 
    1393             ftab = ftab.rename( {az:zdim})  
     1736        if zdim and az : 
     1737            if az != zdim : ftab = ftab.rename( {az:zdim})  
    13941738    return lbc_del (ftab, nperio=nperio, cd_type='F', psgn=psgn) 
    13951739 
    1396 #@numba.jit(forceobj=True) 
    1397 def U2F (utab, nperio=None, psgn=-1.0, zdim='depthf', action='ave') : 
     1740def U2F (utab, nperio=None, psgn=-1.0, zdim=None, action='ave') : 
    13981741    '''Interpolate an array from U grid to F grid i-mean)''' 
    13991742    mmath = __mmath__ (utab) 
    14001743    utab_0 = mmath.where ( np.isnan(utab), 0., utab) 
    14011744    utab_0 = lbc_add (utab_0 , nperio=nperio, cd_type='U', psgn=psgn) 
    1402     iy, ay = __findAxis__ (utab_0, 'y') 
    1403     iz, az = __findAxis__ (utab_0, 'z') 
    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  
     1745    jy, ay = __findAxis__ (utab_0, 'y') 
     1746    kz, az = __findAxis__ (utab_0, 'z') 
     1747    if jy :  
     1748        if action == 'ave'  :    ftab = 0.5 *      (utab_0 + np.roll (utab_0, axis=jy, shift=-1)) 
     1749        if action == 'min'  :    ftab = np.minimum (utab_0 , np.roll (utab_0, axis=jy, shift=-1)) 
     1750        if action == 'max'  :    ftab = np.maximum (utab_0 , np.roll (utab_0, axis=jy, shift=-1)) 
     1751        if action == 'mult' :    ftab =             utab_0 * np.roll (utab_0, axis=jy, shift=-1) 
     1752        ftab = lbc_del (ftab  , nperio=nperio, cd_type='F', psgn=psgn) 
     1753    else : 
     1754        ftab = lbc_del (utab_0, nperio=nperio, cd_type='F', psgn=psgn) 
     1755   
    14101756    if mmath == xr : 
    1411         if ay != None :  
    1412             ftab = ftab.assign_coords({'y':np.arange(ftab.shape[iy])+1.}) 
    1413         if zdim != None and iz != None and az != 'olevel' : 
    1414             ftab = ftab.rename( {az:zdim})  
     1757        if ay :  
     1758            ftab = ftab.assign_coords({'y':np.arange(ftab.shape[jy])+1.}) 
     1759        if zdim and az : 
     1760            if az != zdim : ftab = ftab.rename( {az:zdim})  
    14151761    return ftab 
    14161762 
    1417 #@numba.jit(forceobj=True) 
    1418 def F2T (ftab, nperio=None, psgn=1.0, zdim='deptht', action='ave') : 
     1763def F2T (ftab, nperio=None, psgn=1.0, zdim=None, action='ave') : 
    14191764    '''Interpolate an array on F grid to T grid (i- and j- means)''' 
    14201765    mmath = __mmath__ (ftab) 
    … …  
    14241769    return lbc_del (ttab, nperio=nperio, cd_type='T', psgn=psgn) 
    14251770 
    1426 #@numba.jit(forceobj=True) 
    1427 def T2F (ttab, nperio=None, psgn=1.0, zdim='deptht', action='mean') : 
     1771def T2F (ttab, nperio=None, psgn=1.0, zdim=None, action='mean') : 
    14281772    '''Interpolate an array on T grid to F grid (i- and j- means)''' 
    14291773    mmath = __mmath__ (ttab) 
    14301774    ttab_0 = mmath.where ( np.isnan(ttab), 0., ttab) 
    14311775    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) 
     1776    ftab = T2U (U2F (ttab, nperio=nperio, psgn=psgn, zdim=zdim, action=action), nperio=nperio, psgn=psgn, zdim=zdim, action=action) 
    14331777     
    14341778    return lbc_del (ftab, nperio=nperio, cd_type='F', psgn=psgn) 
    14351779 
    1436 #@numba.jit(forceobj=True) 
    1437 def F2U (ftab, nperio=None, psgn=1.0, zdim='depthu', action='ave') : 
     1780def F2U (ftab, nperio=None, psgn=1.0, zdim=None, action='ave') : 
    14381781    '''Interpolate an array on F grid to FUgrid (i-mean)''' 
    14391782    mmath = __mmath__ (ftab) 
    14401783    ftab_0 = mmath.where ( np.isnan(ftab), 0., ftab) 
    14411784    ftab_0 = lbc_add (ftab_0 , nperio=nperio, cd_type='F', psgn=psgn) 
    1442     iy, ay = __findAxis__ (ftab_0, 'y') 
    1443     iz, az = __findAxis__ (ftab_0, 'z') 
    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      
     1785    jy, ay = __findAxis__ (ftab_0, 'y') 
     1786    kz, az = __findAxis__ (ftab_0, 'z') 
     1787    if jy :  
     1788        if action == 'ave'  : utab = 0.5 *      (ftab_0 + np.roll (ftab_0, axis=jy, shift=-1)) 
     1789        if action == 'min'  : utab = np.minimum (ftab_0 , np.roll (ftab_0, axis=jy, shift=-1)) 
     1790        if action == 'max'  : utab = np.maximum (ftab_0 , np.roll (ftab_0, axis=jy, shift=-1)) 
     1791        if action == 'mult' : utab =             ftab_0 * np.roll (ftab_0, axis=jy, shift=-1) 
     1792        utab = lbc_del (utab  , nperio=nperio, cd_type='U', psgn=psgn) 
     1793    else : 
     1794        utab = lbc_del (ftab_0, nperio=nperio, cd_type='U', psgn=psgn) 
     1795 
    14511796    if mmath == xr : 
    1452         utab = utab.assign_coords({ay:np.arange(ftab.shape[iy])+1.}) 
    1453         if zdim != None and iz != None and az != 'olevel' : 
    1454             utab = utab.rename( {az:zdim})  
     1797        utab = utab.assign_coords({ay:np.arange(ftab.shape[jy])+1.}) 
     1798        if zdim and zz : 
     1799            if az != zdim : utab = utab.rename( {az:zdim})  
    14551800    return utab 
    14561801 
    1457 #@numba.jit(forceobj=True) 
    1458 def F2V (ftab, nperio=None, psgn=1.0, zdim='depthv', action='ave') : 
     1802def F2V (ftab, nperio=None, psgn=1.0, zdim=None, action='ave') : 
    14591803    '''Interpolate an array from F grid to V grid (i-mean)''' 
    14601804    mmath = __mmath__ (ftab) 
    … …  
    14621806    ftab_0 = lbc_add (ftab_0 , nperio=nperio, cd_type='F', psgn=psgn) 
    14631807    ix, ax = __findAxis__ (ftab_0, 'x') 
    1464     iz, az = __findAxis__ (ftab_0, 'z') 
    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) 
     1808    kz, az = __findAxis__ (ftab_0, 'z') 
     1809    if ix :  
     1810        if action == 'ave'  : vtab = 0.5 *      (ftab_0 + np.roll (ftab_0, axis=ix, shift=-1)) 
     1811        if action == 'min'  : vtab = np.minimum (ftab_0 , np.roll (ftab_0, axis=ix, shift=-1)) 
     1812        if action == 'max'  : vtab = np.maximum (ftab_0 , np.roll (ftab_0, axis=ix, shift=-1)) 
     1813        if action == 'mult' : vtab =             ftab_0 * np.roll (ftab_0, axis=ix, shift=-1) 
     1814        vtab = lbc_del (vtab  , nperio=nperio, cd_type='V', psgn=psgn) 
     1815    else :  
     1816        vtab = lbc_del (ftab_0, nperio=nperio, cd_type='V', psgn=psgn) 
     1817 
    14711818    if mmath == xr : 
    14721819        vtab = vtab.assign_coords({ax:np.arange(ftab.shape[ix])+1.}) 
    1473         if zdim != None and iz != None and az != 'olevel' : 
    1474             vtab = vtab.rename( {az:zdim})  
     1820        if zdim and az : 
     1821            if az != zdim : vtab = vtab.rename( {az:zdim})  
    14751822    return vtab 
    14761823 
    1477 #@numba.jit(forceobj=True) 
    1478 def W2T (wtab, zcoord=None, zdim='deptht', sval=np.nan) : 
     1824def W2T (wtab, zcoord=None, zdim=None, sval=np.nan) : 
    14791825    ''' 
    14801826    Interpolate an array on W grid to T grid (k-mean) 
    … …  
    14841830    wtab_0 = mmath.where ( np.isnan(wtab), 0., wtab) 
    14851831 
    1486     iz, az = __findAxis__ (wtab_0, 'z') 
    1487         
    1488     ttab = 0.5 * ( wtab_0 + np.roll (wtab_0, axis=iz, shift=-1) ) 
    1489      
     1832    kz, az = __findAxis__ (wtab_0, 'z') 
     1833 
     1834    if kz :  
     1835        ttab = 0.5 * ( wtab_0 + np.roll (wtab_0, axis=kz, shift=-1) ) 
     1836    else : 
     1837        ttab = wtab_0 
     1838 
    14901839    if mmath == xr : 
    1491         ttab[{az:iz}] = sval 
    1492         if zdim != None and iz != None and az != 'olevel' : 
    1493             ttab = ttab.rename ( {az:zdim} ) 
    1494         try    : ttab = ttab.assign_coords ( {zdim:zcoord} ) 
    1495         except : pass 
     1840        ttab[{az:kz}] = sval 
     1841        if zdim and az : 
     1842            if az != zdim : ttab = ttab.rename ( {az:zdim} ) 
     1843        if zcoord is not None : 
     1844            ttab = ttab.assign_coords ( {zdim:zcoord} ) 
    14961845    else : 
    14971846        ttab[..., -1, :, :] = sval 
    … …  
    14991848    return ttab 
    15001849 
    1501 #@numba.jit(forceobj=True) 
    1502 def T2W (ttab, zcoord=None, zdim='depthw', sval=np.nan, extrap_surf=False) : 
     1850def T2W (ttab, zcoord=None, zdim=None, sval=np.nan, extrap_surf=False) : 
    15031851    '''Interpolate an array from T grid to W grid (k-mean) 
    15041852    sval is the surface value 
    … …  
    15071855    mmath = __mmath__ (ttab) 
    15081856    ttab_0 = mmath.where ( np.isnan(ttab), 0., ttab) 
    1509     iz, az = __findAxis__ (ttab_0, 'z') 
    1510     wtab = 0.5 * ( ttab_0 + np.roll (ttab_0, axis=iz, shift=1) ) 
     1857    kz, az = __findAxis__ (ttab_0, 'z') 
     1858    wtab = 0.5 * ( ttab_0 + np.roll (ttab_0, axis=kz, shift=1) ) 
    15111859 
    15121860    if mmath == xr : 
    … …  
    15181866 
    15191867    if mmath == xr : 
    1520         if zdim != None and iz != None and az != 'olevel' : 
    1521                 wtab = wtab.rename ( {az:zdim}) 
    1522         if zcoord != None : wtab = wtab.assign_coords ( {zdim:zcoord}) 
    1523         else              : ztab = wtab.assign_coords ( {zdim:np.arange(ttab.shape[iz])+1.} ) 
     1868        if zdim and az : 
     1869            if az != zdim : wtab = wtab.rename ( {az:zdim}) 
     1870        if zcoord is not None : 
     1871            wtab = wtab.assign_coords ( {zdim:zcoord}) 
     1872        else : 
     1873            ztab = wtab.assign_coords ( {zdim:np.arange(ttab.shape[kz])+1.} ) 
    15241874    return wtab 
    15251875 
    1526 #@numba.jit(forceobj=True) 
    15271876def fill (ptab, nperio, cd_type='T', npass=1, sval=0.) : 
    15281877    ''' 
    … …  
    15841933    return ztab 
    15851934 
    1586 #@numba.jit(forceobj=True) 
    15871935def correct_uv (u, v, lat) : 
    15881936    ''' 
    1589     Correct a Cartopy bug in Orthographic projection 
     1937    Correct a Cartopy bug in orthographic projection 
    15901938 
    15911939    See https://github.com/SciTools/cartopy/issues/1179 
    … …  
    15931941    The correction is needed with cartopy <= 0.20 
    15941942    It seems that version 0.21 will correct the bug (https://github.com/SciTools/cartopy/pull/1926) 
     1943    Later note : the bug is still present in Cartopy 0.22 
    15951944 
    15961945    Inputs : 
    1597        u, v : eastward/nothward components 
     1946       u, v : eastward/northward components 
    15981947       lat  : latitude of the point (degrees north) 
    15991948 
    … …  
    16081957    return uc, vc 
    16091958 
    1610 def msf (v_e1v_e3v, lat1d, depthw) : 
     1959def norm_uv (u, v) : 
     1960    ''' 
     1961    Return norm of a 2 components vector 
     1962    ''' 
     1963    return np.sqrt (u*u + v*v) 
     1964 
     1965def normalize_uv (u, v) : 
     1966    ''' 
     1967    Normalize 2 components vector 
     1968    ''' 
     1969    uv = norm_uv (u, v) 
     1970    return u/uv, v/uv 
     1971 
     1972def msf (vv, e1v_e3v, lat1d, depthw) : 
    16111973    ''' 
    16121974    Computes the meridonal stream function 
    16131975    First input is meridional_velocity*e1v*e3v 
    16141976    ''' 
    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'}) 
     1977  
     1978    v_e1v_e3v = vv * e1v_e3v 
     1979    v_e1v_e3v.attrs = vv.attrs 
     1980     
     1981    ix, ax = __findAxis__ (v_e1v_e3v, 'x') 
     1982    kz, az = __findAxis__ (v_e1v_e3v, 'z') 
     1983    if az == 'olevel' : new_az = 'olevel' 
     1984    else              : new_az = 'depthw' 
     1985 
     1986    zomsf = -v_e1v_e3v.cumsum ( dim=az, keep_attrs=True).sum (dim=ax, keep_attrs=True)*1.E-6 
     1987    zomsf = zomsf - zomsf.isel ( { az:-1} ) 
     1988     
     1989    jy, ay = __findAxis__ (zomsf, 'y' ) 
     1990    zomsf = zomsf.assign_coords ( {az:depthw.values, ay:lat1d.values}) 
     1991     
     1992    zomsf = zomsf.rename ( {ay:'lat'}) 
     1993    if az != new_az : zomsf = zomsf.rename ( {az:new_az} ) 
     1994    zomsf.attrs['standard_name'] = 'Meridional stream function' 
    16301995    zomsf.attrs['long_name'] = 'Meridional stream function' 
    1631  
    16321996    zomsf.attrs['units'] = 'Sv' 
    1633     zomsf.depthw.attrs=depthw.attrs 
     1997    zomsf[new_az].attrs  = depthw.attrs 
    16341998    zomsf.lat.attrs=lat1d.attrs 
    16351999         
    16362000    return zomsf 
    16372001 
    1638 def bsf (u_e2u_e3u, mask, nperio=None, bsf0=None ) : 
     2002def bsf (uu, e2u_e3u, mask, nperio=None, bsf0=None ) : 
    16392003    ''' 
    16402004    Computes the barotropic stream function 
    16412005    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 : 
     2006    bsf0 is the point with bsf=0  
     2007    (ex: bsf0={'x':3, 'y':120} for orca2,  
     2008         bsf0={'x':5, 'y':300} for eeORCA1 
     2009    ''' 
     2010    u_e2u_e3u       = uu * e2u_e3u 
     2011    u_e2u_e3u.attrs = uu.attrs 
     2012 
     2013    iy, ay = __findAxis__ (u_e2u_e3u, 'y') 
     2014    kz, az = __findAxis__ (u_e2u_e3u, 'z') 
     2015     
     2016    bsf = -u_e2u_e3u.cumsum ( dim=ay, keep_attrs=True ) 
     2017    bsf = bsf.sum (dim=az, keep_attrs=True)*1.E-6 
     2018         
     2019    if bsf0 : 
    16582020        bsf = bsf - bsf.isel (bsf0) 
    16592021        
    16602022    bsf = bsf.where (mask !=0, np.nan) 
    1661     bsf.attrs['long_name'] = 'Barotropic stream function' 
    1662     bsf.attrs['units'] = 'Sv' 
     2023    for attr in uu.attrs : 
     2024        bsf.attrs[attr] = uu.attrs[attr] 
     2025    bsf.attrs['standard_name'] = 'barotropic_stream_function' 
     2026    bsf.attrs['long_name']     = 'Barotropic stream function' 
     2027    bsf.attrs['units']         = 'Sv' 
    16632028    bsf = lbc (bsf, nperio=nperio, cd_type='F') 
    16642029        
    … …  
    16792044     
    16802045    ''' 
    1681  
    1682     if ref != None : 
     2046     
     2047    import f90nml 
     2048    if ref : 
    16832049        if isinstance (ref, str) : nml_ref = f90nml.read (ref) 
    16842050        if isinstance (ref, f90nml.namelist.Namelist) : nml_ref = ref 
    16852051         
    1686     if cfg != None : 
     2052    if cfg : 
    16872053        if isinstance (cfg, str) : nml_cfg = f90nml.read (cfg) 
    16882054        if isinstance (cfg, f90nml.namelist.Namelist) : nml_cfg = cfg 
    … …  
    16932059    list_nml = [] ; list_comment = [] 
    16942060 
    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') 
     2061    if ref : list_nml.append (nml_ref) ; list_comment.append ('ref') 
     2062    if cfg : list_nml.append (nml_cfg) ; list_comment.append ('cfg') 
    16992063 
    17002064    for nml, comment in zip (list_nml, list_comment) : 
    … …  
    17192083    if out == 'xr'   : return xr_namelist 
    17202084 
    1721  
    17222085def fill_closed_seas (imask, nperio=None,  cd_type='T') : 
    17232086    '''Fill closed seas with image processing library 
    … …  
    17302093 
    17312094    return imask_filled 
     2095 
     2096''' 
     2097Sea water state function parameters from NEMO code 
     2098''' 
     2099rdeltaS = 32. ; r1_S0  = 0.875/35.16504 ; r1_T0  = 1./40. ; r1_Z0  = 1.e-4 
     2100 
     2101EOS000 =  8.0189615746e+02 ; EOS100 =  8.6672408165e+02 ; EOS200 = -1.7864682637e+03 ; EOS300 =  2.0375295546e+03 ; EOS400 = -1.2849161071e+03 ; EOS500 =  4.3227585684e+02 ; EOS600 = -6.0579916612e+01 
     2102EOS010 =  2.6010145068e+01 ; EOS110 = -6.5281885265e+01 ; EOS210 =  8.1770425108e+01 ; EOS310 = -5.6888046321e+01 ; EOS410 =  1.7681814114e+01 ; EOS510 = -1.9193502195 
     2103EOS020 = -3.7074170417e+01 ; EOS120 =  6.1548258127e+01 ; EOS220 = -6.0362551501e+01 ; EOS320 =  2.9130021253e+01 ; EOS420 = -5.4723692739     ; EOS030 =  2.1661789529e+01  
     2104EOS130 = -3.3449108469e+01 ; EOS230 =  1.9717078466e+01 ; EOS330 = -3.1742946532 
     2105EOS040 = -8.3627885467     ; EOS140 =  1.1311538584e+01 ; EOS240 = -5.3563304045 
     2106EOS050 =  5.4048723791e-01 ; EOS150 =  4.8169980163e-01 
     2107EOS060 = -1.9083568888e-01 
     2108EOS001 =  1.9681925209e+01 ; EOS101 = -4.2549998214e+01 ; EOS201 =  5.0774768218e+01 ; EOS301 = -3.0938076334e+01 ; EOS401 =   6.6051753097    ; EOS011 = -1.3336301113e+01 
     2109EOS111 = -4.4870114575     ; EOS211 =  5.0042598061     ; EOS311 = -6.5399043664e-01 ; EOS021 =  6.7080479603     ; EOS121 =   3.5063081279 
     2110EOS221 = -1.8795372996     ; EOS031 = -2.4649669534     ; EOS131 = -5.5077101279e-01 ; EOS041 =  5.5927935970e-01 
     2111EOS002 =  2.0660924175     ; EOS102 = -4.9527603989     ; EOS202 =  2.5019633244     ; EOS012 =  2.0564311499     ; EOS112 = -2.1311365518e-01 ; EOS022 = -1.2419983026 
     2112EOS003 = -2.3342758797e-02 ; EOS103 = -1.8507636718e-02 ; EOS013 =  3.7969820455e-01  
     2113 
     2114def rhop ( ptemp, psal ) : 
     2115    ''' 
     2116    Potential density referenced to surface 
     2117    Computation from NEMO code 
     2118    ''' 
     2119    zt  = ptemp * r1_T0                                  # Temperature (°C) 
     2120    zs  = np.sqrt ( np.abs( psal + rdeltaS ) * r1_S0 )   # Square root of salinity (PSS) 
     2121    # 
     2122    prhop = (((((EOS060*zt   \ 
     2123             + EOS150*zs     + EOS050)*zt   \ 
     2124             + (EOS240*zs    + EOS140)*zs + EOS040)*zt   \ 
     2125             + ((EOS330*zs   + EOS230)*zs + EOS130)*zs + EOS030)*zt   \ 
     2126             + (((EOS420*zs  + EOS320)*zs + EOS220)*zs + EOS120)*zs + EOS020)*zt   \ 
     2127             + ((((EOS510*zs + EOS410)*zs + EOS310)*zs + EOS210)*zs + EOS110)*zs + EOS010)*zt   \ 
     2128             + (((((EOS600*zs+ EOS500)*zs + EOS400)*zs + EOS300)*zs + EOS200)*zs + EOS100)*zs + EOS000 
     2129    # 
     2130    return prhop 
     2131 
     2132def rho ( pdep, ptemp, psal ) : 
     2133    ''' 
     2134    In situ density 
     2135    Computation from NEMO code 
     2136    ''' 
     2137    zh  = pdep  * r1_Z0                                  # Depth (m) 
     2138    zt  = ptemp * r1_T0                                  # Temperature (°C) 
     2139    zs  = np.sqrt ( np.abs( psal + rdeltaS ) * r1_S0 )   # Square root salinity (PSS) 
     2140    # 
     2141    zn3 = EOS013*zt + EOS103*zs+EOS003 
     2142    # 
     2143    zn2 = (EOS022*zt + EOS112*zs+EOS012)*zt + (EOS202*zs+EOS102)*zs+EOS002 
     2144    # 
     2145    zn1 = (((EOS041*zt   \ 
     2146         + EOS131*zs   + EOS031)*zt   \ 
     2147         + (EOS221*zs  + EOS121)*zs + EOS021)*zt   \ 
     2148        + ((EOS311*zs  + EOS211)*zs + EOS111)*zs + EOS011)*zt   \ 
     2149        + (((EOS401*zs + EOS301)*zs + EOS201)*zs + EOS101)*zs + EOS001 
     2150    # 
     2151    zn0 = (((((EOS060*zt   \ 
     2152             + EOS150*zs      + EOS050)*zt   \ 
     2153             + (EOS240*zs     + EOS140)*zs + EOS040)*zt   \ 
     2154             + ((EOS330*zs    + EOS230)*zs + EOS130)*zs + EOS030)*zt   \ 
     2155             + (((EOS420*zs   + EOS320)*zs + EOS220)*zs + EOS120)*zs + EOS020)*zt   \ 
     2156             + ((((EOS510*zs  + EOS410)*zs + EOS310)*zs + EOS210)*zs + EOS110)*zs + EOS010)*zt   \ 
     2157             + (((((EOS600*zs + EOS500)*zs + EOS400)*zs + EOS300)*zs + EOS200)*zs + EOS100)*zs + EOS000 
     2158    # 
     2159    prho  = ( ( zn3 * zh + zn2 ) * zh + zn1 ) * zh + zn0 
     2160    # 
     2161    return prho 
    17322162 
    17332163## =========================================================================== 
Note: See TracChangeset for help on using the changeset viewer.