Changeset 6245


Ignore:
Timestamp:
09/20/22 11:40:57 (20 months ago)
Author:
omamce
Message:

O.M. : MOSAIX, new functionnalitie in nemo.py

File:
1 edited

Legend:

Unmodified
Added
Removed

  • TOOLS/MOSAIX/nemo.py

    r6190 r6245  
    22## =========================================================================== 
    33## 
    4 ##  MOSAIX is under CeCILL_V2 licence. See "Licence_CeCILL_V2-en.txt" 
    5 ##  file for an english version of the licence and 
    6 ##  "Licence_CeCILL_V2-fr.txt" for a french version. 
     4##  This software is governed by the CeCILL  license under French law and 
     5##  abiding by the rules of distribution of free software.  You can  use, 
     6##  modify and/ or redistribute the software under the terms of the CeCILL 
     7##  license as circulated by CEA, CNRS and INRIA at the following URL 
     8##  "http://www.cecill.info". 
    79## 
    8 ##  Permission is hereby granted, free of charge, to any person or 
    9 ##  organization obtaining a copy of the software and accompanying 
    10 ##  documentation covered by this license (the "Software") to use, 
    11 ##  reproduce, display, distribute, execute, and transmit the 
    12 ##  Software, and to prepare derivative works of the Software, and to 
    13 ##  permit third-parties to whom the Software is furnished to do so, 
    14 ##  all subject to the following: 
    15 ## 
    16 ##  Warning, to install, configure, run, use any of MOSAIX software or 
    17 ##  to read the associated documentation you'll need at least one (1) 
    18 ##  brain in a reasonably working order. Lack of this implement will 
    19 ##  void any warranties (either express or implied).  Authors assumes 
    20 ##  no responsability for errors, omissions, data loss, or any other 
    21 ##  consequences caused directly or indirectly by the usage of his 
    22 ##  software by incorrectly or partially configured 
     10##  Warning, to install, configure, run, use any of Olivier Marti's 
     11##  software or to read the associated documentation you'll need at least 
     12##  one (1) brain in a reasonably working order. Lack of this implement 
     13##  will void any warranties (either express or implied). 
     14##  O. Marti assumes no responsability for errors, omissions, 
     15##  data loss, or any other consequences caused directly or indirectly by 
     16##  the usage of his software by incorrectly or partially configured 
     17##  personal. 
    2318## 
    2419## =========================================================================== 
    … …  
    3732__HeadURL    = "$HeadURL$" 
    3833 
    39 import sys, numpy as np 
     34import numpy as np 
    4035try    : import xarray as xr 
     36except ImportError : pass 
     37 
     38try    : import f90nml 
     39except : pass 
     40 
     41try : from sklearn.impute import SimpleImputer 
     42except : pass 
     43 
     44try    : import numba 
    4145except : pass 
    4246 
    … …  
    6266tList = [ 't', 'T', 'time'  , ] 
    6367 
    64 ## SVN information 
    65 __Author__   = "$Author$" 
    66 __Date__     = "$Date$" 
    67 __Revision__ = "$Revision$" 
    68 __Id__       = "$Id$" 
    69 __HeadURL    = "$HeadURL$" 
    70  
    71 def __math__ (tab, default=None) : 
    72     math = default 
     68## =========================================================================== 
     69def __mmath__ (tab, default=None) : 
     70    mmath = default 
    7371    try    : 
    74         if type (tab) == xr.core.dataarray.DataArray : math = xr 
     72        if type (tab) == xr.core.dataarray.DataArray : mmath = xr 
    7573    except : 
    7674        pass 
    7775 
    7876    try    : 
    79         if type (tab) == np.ndarray : math = np 
     77        if type (tab) == np.ndarray : mmath = np 
    8078    except : 
    8179        pass 
    8280             
    83     return math 
    84  
    85 def __guessNperio__ (jpj, jpi, nperio=None) : 
     81    return mmath 
     82 
     83 
     84def __guessNperio__ (jpj, jpi, nperio=None, out='nperio') : 
    8685    ''' 
    8786    Tries to guess the value of nperio (periodicity parameter. See NEMO documentation for details) 
    … …  
    9392    ''' 
    9493    if nperio == None : 
     94        nperio = __guessConfig__ (jpj, jpi, nperio=None, out='nperio') 
     95     
     96    return nperio 
     97 
     98def __guessConfig__ (jpj, jpi, nperio=None, config=None, out='nperio') : 
     99    ''' 
     100    Tries to guess the value of nperio (periodicity parameter. See NEMO documentation for details) 
     101     
     102    Inputs 
     103    jpj    : number of latitudes 
     104    jpi    : number of longitudes 
     105    nperio : periodicity parameter 
     106    ''' 
     107    if nperio == None : 
    95108        ## Values for NEMO version < 4.2 
    96         if jpi ==  182 : 
     109        if jpj ==  149 and jpi ==  182 : 
     110            config = 'ORCA2.3' 
    97111            nperio = 4   # ORCA2. We choose legacy orca2. 
    98112            Iperio = 1 ; Jperio = 0 ; NFold = 1 ; NFtype = 'T' 
    99         if jpi ==  362 : # eORCA1. 
     113        if jpj ==  332 and jpi ==  362 : # eORCA1. 
     114            config = 'eORCA1.2' 
    100115            nperio = 6   
    101116            Iperio = 1 ; Jperio = 0 ; NFold = 1 ; NFtype = 'F' 
    102117        if jpi == 1442 :  # ORCA025. 
     118            config = 'ORCA025' 
    103119            nperio = 6  
    104120            Iperio = 1 ; Jperio = 0 ; NFold = 1 ; NFtype = 'F' 
    105         if jpj ==  149 : # ORCA2. We choose legacy orca2. 
    106             nperio = 4 
    107             Iperio = 1 ; Jperio = 0 ; NFold = 1 ; NFtype = 'F' 
    108121        if jpj ==  294 : # ORCA1 
     122            config = 'ORCA1' 
    109123            nperio = 6 
    110             Iperio = 1 ; Jperio = 0 ; NFold = 1 ; NFtype = 'F' 
    111         if jpj ==  332 : # eORCA1. 
    112             nperio = 6   
    113124            Iperio = 1 ; Jperio = 0 ; NFold = 1 ; NFtype = 'F' 
    114125             
    115126        ## Values for NEMO version >= 4.2. No more halo points 
    116         if jpi ==  180 : 
     127        if jpj == 148 and jpi ==  180 : 
     128            config = 'ORCA2.4' 
    117129            nperio = 4.2 # ORCA2. We choose legacy orca2. 
    118130            Iperio = 1 ; Jperio = 0 ; NFold = 1 ; NFtype = 'F' 
    119         if jpi ==  360 : # eORCA1. 
     131        if jpj == 331  and jpi ==  360 : # eORCA1. 
     132            config = 'eORCA1.4' 
    120133            nperio = 6.2 
    121134            Iperio = 1 ; Jperio = 0 ; NFold = 1 ; NFtype = 'F' 
    122135        if jpi == 1440 : # ORCA025. 
     136            config = 'ORCA025' 
    123137            nperio = 6.2 
    124138            Iperio = 1 ; Jperio = 0 ; NFold = 1 ; NFtype = 'F' 
    125  
    126         if jpj ==  148 : nperio = 4.2 # ORCA2. We choose legacy orca2. 
    127         if jpj ==  330 : nperio = 6.2 # eORCA1. 
    128  
     139             
    129140        if nperio == None : 
    130             sys.exit ('in nemo module : nperio not found, and cannot by guessed') 
     141            raise Exception ('in nemo module : nperio not found, and cannot by guessed') 
    131142        else : 
    132143            if nperio in nperio_valid_range : 
    133                 print ('nperio set as {:} (deduced from jpi={:d})'.format (nperio, jpi)) 
     144                print ('nperio set as {:} (deduced from jpj={:d} jpi={:d})'.format (nperio, jpj, jpi)) 
    134145            else :  
    135                 sys.exit  ('nperio set as {:} (deduced from jpi={:d}) : nemo.py is not ready for this value'.format (nperio, jpi)) 
    136              
    137     return nperio 
    138  
     146                raise ValueError ('nperio set as {:} (deduced from jpi={:d}) : nemo.py is not ready for this value'.format (nperio, jpi)) 
     147 
     148    if out == 'nperio' : return nperio 
     149    if out == 'config' : return config 
     150    if out == 'perio'  : return Iperio, Jperio, NFold, NFtype 
     151    if out in ['full', 'all'] : return {'nperio':nperio, 'Iperio':Iperio, 'Jperio':Jperio, 'NFold':NFold, 'NFtype':NFtype} 
     152         
    139153def __guessPoint__ (ptab) : 
    140154    ''' 
    … …  
    149163    ''' 
    150164    gP = None 
    151     math = __math__ (ptab) 
    152     if math == xr : 
     165    mmath = __mmath__ (ptab) 
     166    if mmath == xr : 
    153167        if 'x_c' in ptab.dims and 'y_c' in ptab.dims                        : gP = 'T' 
    154168        if 'x_f' in ptab.dims and 'y_c' in ptab.dims                        : gP = 'U' 
    … …  
    162176              
    163177        if gP == None : 
    164             sys.exit ('in nemo module : cd_type not found, and cannot by guessed') 
     178            raise Exception ('in nemo module : cd_type not found, and cannot by guessed') 
    165179        else : 
    166180            print ('Grid set as', gP, 'deduced from dims ', ptab.dims) 
    167181            return gP 
    168182    else : 
    169          sys.exit ('in nemo module : cd_type not found, input is not an xarray data') 
    170          #return gP 
     183         raise Exception  ('in nemo module : cd_type not found, input is not an xarray data') 
     184 
     185def lbc_diag (nperio) : 
     186    lperio = nperio ; aperio = False 
     187    if nperio == 4.2 : 
     188        lperio = 4 ; aperio = True 
     189    if nperio == 6.2 : 
     190        lperio = 6 ; aperio = True 
     191         
     192    return lperio, aperio  
    171193 
    172194def __findAxis__ (tab, axis='z') : 
    … …  
    174196    Find number and name of the requested axis 
    175197    ''' 
    176     math = __math__ (tab) 
     198    mmath = __mmath__ (tab) 
    177199    ix = None ; ax = None 
    178200 
    … …  
    198220        unList = [ 'second', 'minute', 'hour', 'day', 'month' ] 
    199221     
    200     if math == xr : 
     222    if mmath == xr : 
    201223        for Name in axList : 
    202224            try    : 
    … …  
    222244        
    223245    return ix, ax 
    224                  
     246 
     247#@numba.jit(forceobj=True) 
    225248def fixed_lon (lon, center_lon=0.0) : 
    226249    ''' 
    … …  
    232255    Designed by Phil Pelson. See https://gist.github.com/pelson/79cf31ef324774c97ae7 
    233256    ''' 
    234     math = __math__ (lon) 
     257    mmath = __mmath__ (lon) 
    235258     
    236259    fixed_lon = lon.copy () 
    237260         
    238     fixed_lon = math.where (fixed_lon > center_lon+180., fixed_lon-360.0, fixed_lon) 
    239     fixed_lon = math.where (fixed_lon < center_lon-180., fixed_lon+360.0, fixed_lon) 
     261    fixed_lon = mmath.where (fixed_lon > center_lon+180., fixed_lon-360.0, fixed_lon) 
     262    fixed_lon = mmath.where (fixed_lon < center_lon-180., fixed_lon+360.0, fixed_lon) 
    240263     
    241264    for i, start in enumerate (np.argmax (np.abs (np.diff (fixed_lon, axis=-1)) > 180., axis=-1)) : 
    242         fixed_lon [..., i, start+1:] += 360 
     265        fixed_lon [..., i, start+1:] += 360. 
    243266 
    244267    # Special case for eORCA025 
    … …  
    246269    if fixed_lon.shape [-1] == 1440 : fixed_lon [..., -1, :] = fixed_lon [..., -2, :] 
    247270 
    248     if fixed_lon.min () > center_lon : fixed_lon = fixed_lon-360.0 
     271    if fixed_lon.min () > center_lon : fixed_lon += -360.0 
     272    if fixed_lon.max () < center_lon : fixed_lon +=  360.0 
     273         
     274    if fixed_lon.min () < center_lon-360.0 : fixed_lon +=  360.0 
     275    if fixed_lon.max () > center_lon+360.0 : fixed_lon += -360.0 
    249276                 
    250277    return fixed_lon 
    251278 
     279#@numba.jit(forceobj=True) 
     280def fill_empty (ztab, sval=np.nan, transpose=False) : 
     281    ''' 
     282    Fill values 
     283 
     284    Useful when NEMO has run with no wet points options :  
     285    some parts of the domain, with no ocean points, has no 
     286    lon/lat values 
     287    ''' 
     288    mmath = __mmath__ (ztab) 
     289 
     290    imp = SimpleImputer (missing_values=sval, strategy='mean') 
     291    if transpose : 
     292        imp.fit (ztab.T) 
     293        ptab = imp.transform (ztab.T).T 
     294    else :  
     295        imp.fit (ztab) 
     296        ptab = imp.transform (ztab) 
     297    
     298    if mmath == xr : 
     299        ptab = xr.DataArray (ptab, dims=ztab.dims, coords=ztab.coords) 
     300        ptab.attrs = ztab.attrs 
     301         
     302    return ptab 
     303 
     304#@numba.jit(forceobj=True) 
     305def fill_lonlat (lon, lat, sval=-1) : 
     306    ''' 
     307    Fill longitude/latitude values 
     308 
     309    Useful when NEMO has run with no wet points options :  
     310    some parts of the domain, with no ocean points, as no 
     311    lon/lat values 
     312    ''' 
     313    mmath = __mmath__ (lon) 
     314 
     315    imp = SimpleImputer (missing_values=sval, strategy='mean') 
     316    imp.fit (lon) 
     317    plon = imp.transform (lon) 
     318    imp.fit (lat.T) 
     319    plat = imp.transform (lat.T).T 
     320 
     321    if mmath == xr : 
     322        plon = xr.DataArray (plon, dims=lon.dims, coords=lon.coords) 
     323        plat = xr.DataArray (plat, dims=lat.dims, coords=lat.coords) 
     324        plon.attrs = lon.attrs ; plat.attrs = lat.attrs 
     325         
     326    plon = fixed_lon (plon) 
     327     
     328    return plon, plat 
     329 
     330#@numba.jit(forceobj=True) 
    252331def jeq (lat) : 
    253332    ''' 
    … …  
    256335    lat : latitudes of the grid. At least 2D. 
    257336    ''' 
    258     math = __math__ (lat) 
     337    mmath = __mmath__ (lat) 
    259338    ix, ax = __findAxis__ (lat, 'x') 
    260339    iy, ay = __findAxis__ (lat, 'y') 
    261340 
    262     if math == xr : 
     341    if mmath == xr : 
    263342        jeq = int ( np.mean ( np.argmin (np.abs (np.float64 (lat)), axis=iy) ) ) 
    264343    else :  
    … …  
    266345    return jeq 
    267346 
     347#@numba.jit(forceobj=True) 
    268348def lon1D (lon, lat=None) : 
    269349    ''' 
    … …  
    273353    lat (optionnal) : latitudes of the grid 
    274354    ''' 
     355    mmath = __mmath__ (lon) 
    275356    if np.max (lat) != None : 
    276357        je    = jeq (lat) 
    277         lon1D = lon [..., je, :] 
     358        lon1D = lon.copy() [..., je, :] 
    278359    else : 
    279360        jpj, jpi = lon.shape [-2:] 
    280         lon1D    = lon [..., jpj//3, :] 
     361        lon1D    = lon.copy() [..., jpj//3, :] 
    281362 
    282363    start = np.argmax (np.abs (np.diff (lon1D, axis=-1)) > 180.0, axis=-1) 
    283364    lon1D [..., start+1:] += 360 
     365 
     366    if mmath == xr : 
     367        lon1D.attrs = lon.attrs 
     368        lon1D = lon1D.assign_coords ( {'x':lon1D} ) 
    284369         
    285370    return lon1D 
    286371 
     372#@numba.jit(forceobj=True) 
    287373def latreg (lat, diff=0.1) : 
    288374    ''' 
    289     Returns maximum j index where gridlines are along latitude in the northern hemisphere 
     375    Returns maximum j index where gridlines are along latitudes in the northern hemisphere 
    290376     
    291377    lat : latitudes of the grid (2D) 
    292378    diff [optional] : tolerance 
    293379    ''' 
    294     math = __math__ (lat) 
     380    mmath = __mmath__ (lat) 
    295381    if diff == None : 
    296382        dy   = np.float64 (np.mean (np.abs (lat - np.roll (lat,shift=1,axis=-2, roll_coords=False)))) 
    … …  
    304390    return jreg, latreg 
    305391 
     392#@numba.jit(forceobj=True) 
    306393def lat1D (lat) : 
    307394    ''' 
    … …  
    310397    lat : latitudes of the grid (2D) 
    311398    ''' 
    312     math = __math__ (lat) 
     399    mmath = __mmath__ (lat) 
    313400    jpj, jpi = lat.shape[-2:] 
    314401 
    … …  
    326413        yrange = 90.    -lat_reg 
    327414         
    328     lat1D = math.where (lat_ave<lat_reg, lat_ave, lat_reg + yrange * (np.arange(jpj)-jreg)/(jpj-jreg-1))    
    329          
    330     if math == xr : lat1D.attrs = lat.attrs 
     415    lat1D = mmath.where (lat_ave<lat_reg, lat_ave, lat_reg + yrange * (np.arange(jpj)-jreg)/(jpj-jreg-1))    
     416         
     417    if mmath == xr : 
     418        lat1D.attrs = lat.attrs 
     419        lat1D = lat1D.assign_coords ( {'y':lat1D} ) 
    331420 
    332421    return lat1D 
    333422 
     423#@numba.jit(forceobj=True) 
    334424def latlon1D (lat, lon) : 
    335425    ''' 
    … …  
    340430    return lat1D (lat),  lon1D (lon, lat) 
    341431 
     432##@numba.jit(forceobj=True) 
    342433def mask_lonlat (ptab, x0, x1, y0, y1, lon, lat, sval=np.nan) : 
    343     math = __math__ (ptab) 
     434    mmath = __mmath__ (ptab) 
    344435    try : 
    345436        lon = lon.copy().to_masked_array() 
    … …  
    350441            np.logical_or (np.logical_or (np.logical_and(lon>x0, lon<x1), np.logical_and(lon+360>x0, lon+360<x1)), 
    351442                                      np.logical_and(lon-360>x0, lon-360<x1))) 
    352     tab = math.where (mask, ptab, np.nan) 
     443    tab = mmath.where (mask, ptab, np.nan) 
    353444     
    354445    return tab 
    355          
     446 
     447#@numba.jit(forceobj=True)       
    356448def extend (tab, Lon=False, jplus=25, jpi=None, nperio=4) : 
    357449    ''' 
    … …  
    366458     
    367459    ''' 
    368     math = __math__ (tab) 
     460    mmath = __mmath__ (tab) 
    369461     
    370462    if tab.shape[-1] == 1 : extend = tab 
    … …  
    386478                istart = 1 ; le=jpi-2 ; la=1  # Perfect, except at the pole that should be masked by lbc_plot 
    387479            
    388             if math == xr : 
     480            if mmath == xr : 
    389481                extend = np.concatenate ((tab.values[..., istart   :istart+le+1    ] + xplus, 
    390482                                          tab.values[..., istart+la:istart+la+jplus]         ), axis=-1) 
    … …  
    402494def orca2reg (ff, lat_name='nav_lat', lon_name='nav_lon', y_name='y', x_name='x') : 
    403495    ''' 
    404     Assign ORCA dataset on a regular grid. 
     496    Assign an ORCA dataset on a regular grid. 
    405497    For use in the tropical region. 
    406  
     498     
    407499    Inputs :  
    408500      ff : xarray dataset 
    409501      lat_name, lon_name : name of latitude and longitude 2D field in ff 
    410502      y_name, x_name     : namex of dimensions in ff 
    411  
    412     Returns : xarray dataset with rectangular grid. Incorrect above 20°N 
     503       
     504      Returns : xarray dataset with rectangular grid. Incorrect above 20°N 
    413505    ''' 
    414506    # Compute 1D longitude and latitude 
    415507    (lat, lon) = latlon1D (ff[lat_name], ff[lon_name]) 
     508 
    416509    # Assign lon and lat as dimensions of the dataset 
    417510    if y_name in ff.dims :  
    … …  
    450543     
    451544    if nperio not in nperio_valid_range : 
    452         print ('nperio=', nperio, ' is not in the valid range', nperio_valid_range) 
    453         sys.exit () 
     545        raise Exception ('nperio=', nperio, ' is not in the valid range', nperio_valid_range) 
    454546 
    455547    return jpj, jpi, nperio 
    456548         
    457          
     549#@numba.jit(forceobj=True) 
    458550def lbc (ptab, nperio=None, cd_type='T', psgn=1.0, nemo_4U_bug=False) : 
    459551    ''' 
    … …  
    468560    jpj, jpi, nperio = lbc_init (ptab, nperio) 
    469561    psgn   = ptab.dtype.type (psgn) 
    470     math = __math__ (ptab) 
    471      
    472     if math == xr : ztab = ptab.values.copy () 
    473     else          : ztab = ptab.copy () 
     562    mmath = __mmath__ (ptab) 
     563     
     564    if mmath == xr : ztab = ptab.values.copy () 
     565    else           : ztab = ptab.copy () 
     566         
    474567    # 
    475568    #> East-West boundary conditions 
    … …  
    548641        ztab [..., :, -1] = ztab [..., :,  1] 
    549642 
    550     if math == xr : 
     643    if mmath == xr : 
    551644        ztab = xr.DataArray ( ztab, dims=ptab.dims, coords=ptab.coords) 
     645        ztab.attrs = ptab.attrs 
    552646         
    553647    return ztab 
    554648 
     649#@numba.jit(forceobj=True) 
    555650def lbc_mask (ptab, nperio=None, cd_type='T', sval=np.nan) : 
    556651    # 
    … …  
    632727    return ztab 
    633728 
     729#@numba.jit(forceobj=True) 
    634730def lbc_plot (ptab, nperio=None, cd_type='T', psgn=1.0, sval=np.nan) : 
    635731    ''' 
    … …  
    707803    return ztab 
    708804 
    709 def lbc_add (ptab, nperio=None, cd_type =None) : 
     805#@numba.jit(forceobj=True) 
     806def lbc_add (ptab, nperio=None, cd_type=None, psgn=1, sval=None) : 
    710807    ''' 
    711808    Handle NEMO domain changes between NEMO 4.0 to NEMO 4.2 
    712       Peridodicity halo has been removes 
     809      Peridodicity halo has been removed 
    713810    This routine adds the halos if needed 
    714811 
    … …  
    719816    See NEMO documentation for further details 
    720817    ''' 
    721  
     818    mmath = __mmath__ (ptab)  
    722819    jpj, jpi, nperio = lbc_init (ptab, nperio) 
    723820 
    … …  
    725822 
    726823    if nperio == 4.2 or nperio == 6.2 : 
    727         math = __math__ (ptab)  
    728824       
    729825        ext_shape = t_shape 
    … …  
    731827        ext_shape[-2] = ext_shape[-2] + 1 
    732828 
    733         if math == xr : ptab_ext = xr.DataArray (np.empty (ext_shape), dims=ptab.dims)  
    734         else          : ptab_ext =               np.empty (ext_shape) 
     829        if mmath == xr : 
     830            ptab_ext = xr.DataArray (np.zeros (ext_shape), dims=ptab.dims)  
     831            ptab_ext.values[..., :-1, 1:-1] = ptab.values.copy () 
     832        else           : 
     833            ptab_ext =               np.zeros (ext_shape) 
     834            ptab_ext[..., :-1, 1:-1] = ptab.copy () 
    735835             
    736         ptab_ext[..., :-1, 1:-1] = ptab 
    737          
    738         if nperio == 4.2 : ptab_ext = lbc (ptab_ext, 4, cd_type) 
    739         if nperio == 6.2 : ptab_ext = lbc (ptab_ext, 6, cd_type) 
    740  
    741         if math == xr : 
    742             for attr in ptab.attrs : 
    743                 ptab_ext.attrs [attr] = ptab.attrs [attr] 
    744  
    745     else : ptab_ext = ptab 
    746          
    747     return ptab 
    748  
    749 def lbc_del (ptab, nperio=None) : 
     836        #if sval != None :  ptab_ext[..., 0, :] = sval 
     837         
     838        if nperio == 4.2 : ptab_ext = lbc (ptab_ext, nperio=4, cd_type=cd_type, psgn=psgn) 
     839        if nperio == 6.2 : ptab_ext = lbc (ptab_ext, nperio=6, cd_type=cd_type, psgn=psgn) 
     840              
     841        if mmath == xr : 
     842            ptab_ext.attrs = ptab.attrs 
     843 
     844    else : ptab_ext = lbc (ptab, nperio=nperio, cd_type=cd_type, psgn=psgn) 
     845         
     846    return ptab_ext 
     847 
     848def lbc_del (ptab, nperio=None, cd_type='T', psgn=1) : 
    750849    ''' 
    751850    Handle NEMO domain changes between NEMO 4.0 to NEMO 4.2 
    752       Peridodicity halo has been removes 
     851      Periodicity halo has been removed 
    753852    This routine removes the halos if needed 
    754853 
    … …  
    762861    jpj, jpi, nperio = lbc_init (ptab, nperio) 
    763862 
    764     if nperio == 4 or nperio == 6 : 
    765         return ptab[..., :-1, 1:-1] 
     863    if nperio == 4.2 or nperio == 6.2 : 
     864        return lbc (ptab[..., :-1, 1:-1], nperio=nperio, cd_type=cd_type, psgn=psgn) 
    766865    else : 
    767866        return ptab 
    768867 
     868#@numba.jit(forceobj=True) 
    769869def lbc_index (jj, ii, jpj, jpi, nperio=None, cd_type='T') : 
    770870    ''' 
    … …  
    784884    jy = jj + 1 ; ix = ii + 1 
    785885 
    786     math = __math__ (jj) 
    787     if math == None : math=np 
     886    mmath = __mmath__ (jj) 
     887    if mmath == None : mmath=np 
    788888 
    789889    # 
    … …  
    792892    if nperio in [1, 4, 6] : 
    793893        #... cyclic 
    794         ix = math.where (ix==jpi, 2   , ix) 
    795         ix = math.where (ix== 1 ,jpi-1, ix) 
     894        ix = mmath.where (ix==jpi, 2   , ix) 
     895        ix = mmath.where (ix== 1 ,jpi-1, ix) 
    796896 
    797897    # 
    798898    def modIJ (cond, jy_new, ix_new) : 
    799         jy_r = math.where (cond, jy_new, jy) 
    800         ix_r = math.where (cond, ix_new, ix) 
     899        jy_r = mmath.where (cond, jy_new, jy) 
     900        ix_r = mmath.where (cond, ix_new, ix) 
    801901        return jy_r, ix_r 
    802902    # 
    … …  
    9071007    else                         : jpj = len(lat_grid) ; jpi=len(lon_grid) 
    9081008 
    909     math = __math__ (lat_grid) 
     1009    mmath = __mmath__ (lat_grid) 
    9101010         
    9111011    # Compute distance from the point to all grid points (in radian) 
    … …  
    9431043def clo_lon (lon, lon0) : 
    9441044    '''Choose closest to lon0 longitude, adding or substacting 360° if needed''' 
    945     math = __math__ (lon, np) 
     1045    mmath = __mmath__ (lon, np) 
    9461046         
    9471047    clo_lon = lon 
    948     clo_lon = math.where (clo_lon > lon0 + 180., clo_lon-360., clo_lon) 
    949     clo_lon = math.where (clo_lon < lon0 - 180., clo_lon+360., clo_lon) 
    950     clo_lon = math.where (clo_lon > lon0 + 180., clo_lon-360., clo_lon) 
    951     clo_lon = math.where (clo_lon < lon0 - 180., clo_lon+360., clo_lon) 
     1048    clo_lon = mmath.where (clo_lon > lon0 + 180., clo_lon-360., clo_lon) 
     1049    clo_lon = mmath.where (clo_lon < lon0 - 180., clo_lon+360., clo_lon) 
     1050    clo_lon = mmath.where (clo_lon > lon0 + 180., clo_lon-360., clo_lon) 
     1051    clo_lon = mmath.where (clo_lon < lon0 - 180., clo_lon+360., clo_lon) 
    9521052    if clo_lon.shape == () : clo_lon = clo_lon.item () 
    9531053    return clo_lon 
    … …  
    9551055def angle_full (glamt, gphit, glamu, gphiu, glamv, gphiv, glamf, gphif, nperio=None) : 
    9561056    '''Compute sinus and cosinus of model line direction with respect to east''' 
    957     math = __math__ (glamt) 
     1057    mmath = __mmath__ (glamt) 
     1058 
     1059    zlamt = lbc_add (glamt, nperio, 'T', 1.) 
     1060    zphit = lbc_add (gphit, nperio, 'T', 1.) 
     1061    zlamu = lbc_add (glamu, nperio, 'U', 1.) 
     1062    zphiu = lbc_add (gphiu, nperio, 'U', 1.) 
     1063    zlamv = lbc_add (glamv, nperio, 'V', 1.) 
     1064    zphiv = lbc_add (gphiv, nperio, 'V', 1.) 
     1065    zlamf = lbc_add (glamf, nperio, 'F', 1.) 
     1066    zphif = lbc_add (gphif, nperio, 'F', 1.) 
    9581067     
    9591068    # north pole direction & modulous (at T-point) 
    960     zxnpt = 0. - 2.0 * np.cos (rad*glamt) * np.tan (rpi/4.0 - rad*gphit/2.0) 
    961     zynpt = 0. - 2.0 * np.sin (rad*glamt) * np.tan (rpi/4.0 - rad*gphit/2.0) 
     1069    zxnpt = 0. - 2.0 * np.cos (rad*zlamt) * np.tan (rpi/4.0 - rad*zphit/2.0) 
     1070    zynpt = 0. - 2.0 * np.sin (rad*zlamt) * np.tan (rpi/4.0 - rad*zphit/2.0) 
    9621071    znnpt = zxnpt*zxnpt + zynpt*zynpt 
    9631072     
    9641073    # north pole direction & modulous (at U-point) 
    965     zxnpu = 0. - 2.0 * np.cos (rad*glamu) * np.tan (rpi/4.0 - rad*gphiu/2.0) 
    966     zynpu = 0. - 2.0 * np.sin (rad*glamu) * np.tan (rpi/4.0 - rad*gphiu/2.0) 
     1074    zxnpu = 0. - 2.0 * np.cos (rad*zlamu) * np.tan (rpi/4.0 - rad*zphiu/2.0) 
     1075    zynpu = 0. - 2.0 * np.sin (rad*zlamu) * np.tan (rpi/4.0 - rad*zphiu/2.0) 
    9671076    znnpu = zxnpu*zxnpu + zynpu*zynpu 
    9681077     
    9691078    # north pole direction & modulous (at V-point) 
    970     zxnpv = 0. - 2.0 * np.cos (rad*glamv) * np.tan (rpi/4.0 - rad*gphiv/2.0) 
    971     zynpv = 0. - 2.0 * np.sin (rad*glamv) * np.tan (rpi/4.0 - rad*gphiv/2.0) 
     1079    zxnpv = 0. - 2.0 * np.cos (rad*zlamv) * np.tan (rpi/4.0 - rad*zphiv/2.0) 
     1080    zynpv = 0. - 2.0 * np.sin (rad*zlamv) * np.tan (rpi/4.0 - rad*zphiv/2.0) 
    9721081    znnpv = zxnpv*zxnpv + zynpv*zynpv 
    9731082 
    9741083    # north pole direction & modulous (at F-point) 
    975     zxnpf = 0. - 2.0 * np.cos( rad*glamf ) * np.tan ( rpi/4. - rad*gphif/2. ) 
    976     zynpf = 0. - 2.0 * np.sin( rad*glamf ) * np.tan ( rpi/4. - rad*gphif/2. ) 
     1084    zxnpf = 0. - 2.0 * np.cos( rad*zlamf ) * np.tan ( rpi/4. - rad*zphif/2. ) 
     1085    zynpf = 0. - 2.0 * np.sin( rad*zlamf ) * np.tan ( rpi/4. - rad*zphif/2. ) 
    9771086    znnpf = zxnpf*zxnpf + zynpf*zynpf 
    9781087 
    9791088    # j-direction: v-point segment direction (around T-point) 
    980     zlam = glamv   
    981     zphi = gphiv 
    982     zlan = np.roll ( glamv, axis=-2, shift=1)  # glamv (ji,jj-1) 
    983     zphh = np.roll ( gphiv, axis=-2, shift=1)  # gphiv (ji,jj-1) 
     1089    zlam = zlamv   
     1090    zphi = zphiv 
     1091    zlan = np.roll ( zlamv, axis=-2, shift=1)  # glamv (ji,jj-1) 
     1092    zphh = np.roll ( zphiv, axis=-2, shift=1)  # gphiv (ji,jj-1) 
    9841093    zxvvt =  2.0 * np.cos ( rad*zlam ) * np.tan ( rpi/4. - rad*zphi/2. )   \ 
    9851094          -  2.0 * np.cos ( rad*zlan ) * np.tan ( rpi/4. - rad*zphh/2. ) 
    … …  
    9891098 
    9901099    # j-direction: f-point segment direction (around u-point) 
    991     zlam = glamf 
    992     zphi = gphif 
    993     zlan = np.roll (glamf, axis=-2, shift=1) # glamf (ji,jj-1) 
    994     zphh = np.roll (gphif, axis=-2, shift=1) # gphif (ji,jj-1) 
     1100    zlam = zlamf 
     1101    zphi = zphif 
     1102    zlan = np.roll (zlamf, axis=-2, shift=1) # glamf (ji,jj-1) 
     1103    zphh = np.roll (zphif, axis=-2, shift=1) # gphif (ji,jj-1) 
    9951104    zxffu =  2.0 * np.cos ( rad*zlam ) * np.tan ( rpi/4. - rad*zphi/2. )   \ 
    9961105          -  2.0 * np.cos ( rad*zlan ) * np.tan ( rpi/4. - rad*zphh/2. ) 
    … …  
    10001109 
    10011110    # i-direction: f-point segment direction (around v-point) 
    1002     zlam = glamf   
    1003     zphi = gphif 
    1004     zlan = np.roll (glamf, axis=-1, shift=1) # glamf (ji-1,jj) 
    1005     zphh = np.roll (gphif, axis=-1, shift=1) # gphif (ji-1,jj) 
     1111    zlam = zlamf   
     1112    zphi = zphif 
     1113    zlan = np.roll (zlamf, axis=-1, shift=1) # glamf (ji-1,jj) 
     1114    zphh = np.roll (zphif, axis=-1, shift=1) # gphif (ji-1,jj) 
    10061115    zxffv =  2.0 * np.cos ( rad*zlam ) * np.tan ( rpi/4. - rad*zphi/2. )   \ 
    10071116          -  2.0 * np.cos ( rad*zlan ) * np.tan ( rpi/4. - rad*zphh/2. ) 
    … …  
    10111120 
    10121121    # j-direction: u-point segment direction (around f-point) 
    1013     zlam = np.roll (glamu, axis=-2, shift=-1) # glamu (ji,jj+1)  
    1014     zphi = np.roll (gphiu, axis=-2, shift=-1) # gphiu (ji,jj+1) 
    1015     zlan = glamu 
    1016     zphh = gphiu 
     1122    zlam = np.roll (zlamu, axis=-2, shift=-1) # glamu (ji,jj+1)  
     1123    zphi = np.roll (zphiu, axis=-2, shift=-1) # gphiu (ji,jj+1) 
     1124    zlan = zlamu 
     1125    zphh = zphiu 
    10171126    zxuuf =  2. * np.cos ( rad*zlam ) * np.tan ( rpi/4. - rad*zphi/2. )   \ 
    10181127          -  2. * np.cos ( rad*zlan ) * np.tan ( rpi/4. - rad*zphh/2. ) 
    … …  
    10351144    gcosv =-( zxnpv*zyffv - zynpv*zxffv ) / znffv  # (caution, rotation of 90 degres) 
    10361145     
    1037     gsint = lbc (gsint, cd_type='T', nperio=nperio, psgn=-1.) 
    1038     gcost = lbc (gcost, cd_type='T', nperio=nperio, psgn=-1.) 
    1039     gsinu = lbc (gsinu, cd_type='U', nperio=nperio, psgn=-1.) 
    1040     gcosu = lbc (gcosu, cd_type='U', nperio=nperio, psgn=-1.) 
    1041     gsinv = lbc (gsinv, cd_type='V', nperio=nperio, psgn=-1.) 
    1042     gcosv = lbc (gcosv, cd_type='V', nperio=nperio, psgn=-1.) 
    1043     gsinf = lbc (gsinf, cd_type='F', nperio=nperio, psgn=-1.) 
    1044     gcosf = lbc (gcosf, cd_type='F', nperio=nperio, psgn=-1.) 
    1045  
    1046     if math == xr : 
     1146    #gsint = lbc (gsint, cd_type='T', nperio=nperio, psgn=-1.) 
     1147    #gcost = lbc (gcost, cd_type='T', nperio=nperio, psgn=-1.) 
     1148    #gsinu = lbc (gsinu, cd_type='U', nperio=nperio, psgn=-1.) 
     1149    #gcosu = lbc (gcosu, cd_type='U', nperio=nperio, psgn=-1.) 
     1150    #gsinv = lbc (gsinv, cd_type='V', nperio=nperio, psgn=-1.) 
     1151    #gcosv = lbc (gcosv, cd_type='V', nperio=nperio, psgn=-1.) 
     1152    #gsinf = lbc (gsinf, cd_type='F', nperio=nperio, psgn=-1.) 
     1153    #gcosf = lbc (gcosf, cd_type='F', nperio=nperio, psgn=-1.) 
     1154 
     1155    gsint = lbc_del (gsint, cd_type='T', nperio=nperio, psgn=-1.) 
     1156    gcost = lbc_del (gcost, cd_type='T', nperio=nperio, psgn=-1.) 
     1157    gsinu = lbc_del (gsinu, cd_type='U', nperio=nperio, psgn=-1.) 
     1158    gcosu = lbc_del (gcosu, cd_type='U', nperio=nperio, psgn=-1.) 
     1159    gsinv = lbc_del (gsinv, cd_type='V', nperio=nperio, psgn=-1.) 
     1160    gcosv = lbc_del (gcosv, cd_type='V', nperio=nperio, psgn=-1.) 
     1161    gsinf = lbc_del (gsinf, cd_type='F', nperio=nperio, psgn=-1.) 
     1162    gcosf = lbc_del (gcosf, cd_type='F', nperio=nperio, psgn=-1.) 
     1163 
     1164    if mmath == xr : 
    10471165        gsint = gsint.assign_coords ( glamt.coords ) 
    10481166        gcost = gcost.assign_coords ( glamt.coords ) 
    … …  
    10581176def angle (glam, gphi, nperio, cd_type='T') : 
    10591177    '''Compute sinus and cosinus of model line direction with respect to east''' 
    1060     math = __math__ (glam)  
     1178    mmath = __mmath__ (glam) 
     1179 
     1180    zlam = lbc_add (glam, nperio, cd_type, 1.) 
     1181    zphi = lbc_add (gphi, nperio, cd_type, 1.) 
     1182     
    10611183    # north pole direction & modulous 
    1062     zxnp = 0. - 2.0 * np.cos (rad*glam) * np.tan (rpi/4.0 - rad*gphi/2.0) 
    1063     zynp = 0. - 2.0 * np.sin (rad*glam) * np.tan (rpi/4.0 - rad*gphi/2.0) 
     1184    zxnp = 0. - 2.0 * np.cos (rad*zlam) * np.tan (rpi/4.0 - rad*zphi/2.0) 
     1185    zynp = 0. - 2.0 * np.sin (rad*zlam) * np.tan (rpi/4.0 - rad*zphi/2.0) 
    10641186    znnp = zxnp*zxnp + zynp*zynp 
    10651187 
    10661188    # j-direction: segment direction (around point) 
    1067     zlan_n = np.roll (glam, axis=-2, shift=-1) # glam [jj+1, ji] 
    1068     zphh_n = np.roll (gphi, axis=-2, shift=-1) # gphi [jj+1, ji] 
    1069     zlan_s = np.roll (glam, axis=-2, shift= 1) # glam [jj-1, ji] 
    1070     zphh_s = np.roll (gphi, axis=-2, shift= 1) # gphi [jj-1, ji] 
     1189    zlan_n = np.roll (zlam, axis=-2, shift=-1) # glam [jj+1, ji] 
     1190    zphh_n = np.roll (zphi, axis=-2, shift=-1) # gphi [jj+1, ji] 
     1191    zlan_s = np.roll (zlam, axis=-2, shift= 1) # glam [jj-1, ji] 
     1192    zphh_s = np.roll (zphi, axis=-2, shift= 1) # gphi [jj-1, ji] 
    10711193     
    10721194    zxff = 2.0 * np.cos (rad*zlan_n) * np.tan (rpi/4.0 - rad*zphh_n/2.0) \ 
    … …  
    10791201    gcos = (zxnp*zxff + zynp*zyff) / znff 
    10801202 
    1081     gsin = lbc (gsin, cd_type=cd_type, nperio=nperio, psgn=-1.) 
    1082     gcos = lbc (gcos, cd_type=cd_type, nperio=nperio, psgn=-1.) 
    1083  
    1084     if math == xr : 
     1203    gsin = lbc_del (gsin, cd_type=cd_type, nperio=nperio, psgn=-1.) 
     1204    gcos = lbc_del (gcos, cd_type=cd_type, nperio=nperio, psgn=-1.) 
     1205 
     1206    if mmath == xr : 
    10851207        gsin = gsin.assign_coords ( glam.coords ) 
    10861208        gcos = gcos.assign_coords ( glam.coords ) 
    … …  
    11241246    east-north components on the T grid point    
    11251247    ''' 
    1126     math = __math__ (uo) 
     1248    mmath = __mmath__ (uo) 
    11271249 
    11281250    ut = U2T (uo, nperio=nperio, psgn=-1.0, zdim=zdim) 
    … …  
    11391261def rot_uv2enF ( uo, vo, gsinf, gcosf, nperio, zdim='deptht' ) : 
    11401262    ''' 
    1141     ** Purpose :   Rotate the Repere: Change vector componantes from 
     1263    ** Purpose : Rotate the Repere: Change vector componantes from 
    11421264    stretched coordinates grid --> geographic grid 
    11431265    uo is on the U grid point, vo is on the V grid point 
    11441266    east-north components on the T grid point    
    11451267    ''' 
    1146     math = __math__ (uo) 
     1268    mmath = __mmath__ (uo) 
    11471269 
    11481270    uf = U2F (uo, nperio=nperio, psgn=-1.0, zdim=zdim) 
    … …  
    11561278     
    11571279    return u_e, v_n 
    1158               
    1159 def U2T (utab, nperio=None, psgn=-1.0, zdim='deptht') : 
     1280 
     1281#@numba.jit(forceobj=True) 
     1282def U2T (utab, nperio=None, psgn=-1.0, zdim='deptht', action='ave') : 
    11601283    '''Interpolate an array from U grid to T grid i-mean)''' 
    1161     math = __math__ (utab) 
    1162     utab_0 = math.where ( np.isnan(utab), 0., utab) 
    1163     utab_0 = lbc (utab_0 , nperio=nperio, cd_type='U', psgn=psgn) 
     1284    mmath = __mmath__ (utab) 
     1285    utab_0 = mmath.where ( np.isnan(utab), 0., utab) 
     1286    lperio, aperio = lbc_diag (nperio) 
     1287    utab_0 = lbc_add (utab_0, nperio=nperio, cd_type='U', psgn=psgn) 
    11641288    ix, ax = __findAxis__ (utab_0, 'x') 
    11651289    iz, az = __findAxis__ (utab_0, 'z') 
    1166     ttab =  lbc ( 0.5 * (utab_0 + np.roll (utab_0, axis=ix, shift=1)), cd_type='T', nperio=nperio, psgn=psgn) 
    1167      
    1168     if math == xr : 
     1290    if action == 'ave'  : ttab = 0.5 *      (utab_0 + np.roll (utab_0, axis=ix, shift=1)) 
     1291    if action == 'min'  : ttab = np.minimum (utab_0 , np.roll (utab_0, axis=ix, shift=1)) 
     1292    if action == 'max'  : ttab = np.maximum (utab_0 , np.roll (utab_0, axis=ix, shift=1)) 
     1293    if action == 'mult' : ttab =             utab_0 * np.roll (utab_0, axis=ix, shift=1) 
     1294    ttab = lbc_del (ttab, nperio=nperio, cd_type='T', psgn=psgn) 
     1295     
     1296    if mmath == xr : 
    11691297        if ax != None : 
    11701298            ttab = ttab.assign_coords({ax:np.arange (ttab.shape[ix])+1.}) 
    … …  
    11731301    return ttab 
    11741302 
    1175 def V2T (vtab, nperio=None, psgn=-1.0, zdim='deptht') : 
     1303#@numba.jit(forceobj=True) 
     1304def V2T (vtab, nperio=None, psgn=-1.0, zdim='deptht', action='ave') : 
    11761305    '''Interpolate an array from V grid to T grid (j-mean)''' 
    1177     math = __math__ (vtab) 
    1178     vtab_0 = math.where ( np.isnan(vtab), 0., vtab) 
    1179     vtab_0 = lbc (vtab_0 , nperio=nperio, cd_type='V', psgn=psgn) 
     1306    mmath = __mmath__ (vtab) 
     1307    lperio, aperio = lbc_diag (nperio) 
     1308    vtab_0 = mmath.where ( np.isnan(vtab), 0., vtab) 
     1309    vtab_0 = lbc_add (vtab_0, nperio=nperio, cd_type='V', psgn=psgn) 
    11801310    iy, ay = __findAxis__ (vtab_0, 'y') 
    11811311    iz, az = __findAxis__ (vtab_0, 'z') 
    1182     ttab = lbc ( 0.5 * (vtab_0 + np.roll (vtab_0, axis=iy, shift=1)), cd_type='T', nperio=nperio, psgn=psgn) 
    1183                       
    1184     if math == xr : 
     1312    if action == 'ave'  : ttab = 0.5 *      (vtab_0 + np.roll (vtab_0, axis=iy, shift=1)) 
     1313    if action == 'min'  : ttab = np.minimum (vtab_0 , np.roll (vtab_0, axis=iy, shift=1)) 
     1314    if action == 'max'  : ttab = np.maximum (vtab_0 , np.roll (vtab_0, axis=iy, shift=1)) 
     1315    if action == 'mult' : ttab =             vtab_0 * np.roll (vtab_0, axis=iy, shift=1) 
     1316    ttab = lbc_del (ttab, nperio=nperio, cd_type='T', psgn=psgn) 
     1317    if mmath == xr : 
    11851318        if ay !=None :  
    11861319            ttab = ttab.assign_coords({ay:np.arange(ttab.shape[iy])+1.}) 
    … …  
    11891322    return ttab 
    11901323 
    1191 def F2T (ftab, nperio=None, psgn=1.0, zdim='depthf') : 
     1324#@numba.jit(forceobj=True) 
     1325def F2T (ftab, nperio=None, psgn=1.0, zdim='depthf', action='ave') : 
    11921326    '''Interpolate an array from F grid to T grid (i- and j- means)''' 
    1193     math = __math__ (ftab) 
    1194     ftab_0 = math.where ( np.isnan(ftab), 0., ftab) 
    1195     ftab_0 = lbc (ftab_0 , nperio=nperio, cd_type='F', psgn=psgn) 
    1196     ttab = V2T(F2V(ftab_0, nperio=nperio, psgn=psgn, zdim=zdim), nperio=nperio, psgn=psgn, zdim=zdim) 
    1197     return ttab 
    1198  
    1199 def T2U (ttab, nperio=None, psgn=1.0, zdim='depthu') : 
     1327    mmath = __mmath__ (ftab) 
     1328    ftab_0 = mmath.where ( np.isnan(ftab), 0., ftab) 
     1329    ftab_0 = lbc_add (ftab_0 , nperio=nperio, cd_type='F', psgn=psgn) 
     1330    ttab = V2T(F2V(ftab_0, nperio=nperio, psgn=psgn, zdim=zdim, action=action), nperio=nperio, psgn=psgn, zdim=zdim, action=action) 
     1331    return lbc_del (ttab, nperio=nperio, cd_type='T', psgn=psgn) 
     1332 
     1333#@numba.jit(forceobj=True) 
     1334def T2U (ttab, nperio=None, psgn=1.0, zdim='depthu', action='ave') : 
    12001335    '''Interpolate an array from T grid to U grid (i-mean)''' 
    1201     math = __math__ (ttab) 
    1202     ttab_0 = math.where ( np.isnan(ttab), 0., ttab) 
    1203     ttab_0 = lbc (ttab_0 , nperio=nperio, cd_type='T', psgn=psgn) 
     1336    mmath = __mmath__ (ttab) 
     1337    ttab_0 = mmath.where ( np.isnan(ttab), 0., ttab) 
     1338    ttab_0 = lbc_add (ttab_0 , nperio=nperio, cd_type='T', psgn=psgn) 
    12041339    ix, ax = __findAxis__ (ttab_0, 'x') 
    12051340    iz, az = __findAxis__ (ttab_0, 'z') 
    1206     utab = lbc ( 0.5 * (ttab_0 + np.roll (ttab_0, axis=ix, shift=-1)), cd_type='U', nperio=nperio, psgn=psgn) 
    1207  
    1208     if math == xr :     
     1341    if action == 'ave'  : utab = 0.5 *      (ttab_0 + np.roll (ttab_0, axis=ix, shift=-1)) 
     1342    if action == 'min'  : utab = np.minimum (ttab_0 , np.roll (ttab_0, axis=ix, shift=-1)) 
     1343    if action == 'max'  : utab = np.maximum (ttab_0 , np.roll (ttab_0, axis=ix, shift=-1)) 
     1344    if action == 'mult' : utab =             ttab_0 * np.roll (ttab_0, axis=ix, shift=-1) 
     1345    utab = lbc_del (utab, nperio=nperio, cd_type='U', psgn=psgn) 
     1346 
     1347    if mmath == xr :     
    12091348        if ax != None :  
    12101349            utab = ttab.assign_coords({ax:np.arange(utab.shape[ix])+1.}) 
    … …  
    12131352    return utab 
    12141353 
    1215 def T2V (ttab, nperio=None, psgn=1.0, zdim='depthv') : 
     1354#@numba.jit(forceobj=True) 
     1355def T2V (ttab, nperio=None, psgn=1.0, zdim='depthv', action='ave') : 
    12161356    '''Interpolate an array from T grid to V grid (j-mean)''' 
    1217     math = __math__ (ttab) 
    1218     ttab_0 = math.where ( np.isnan(ttab), 0., ttab) 
    1219     ttab_0 = lbc (ttab_0 , nperio=nperio, cd_type='T', psgn=psgn) 
     1357    mmath = __mmath__ (ttab) 
     1358    ttab_0 = mmath.where ( np.isnan(ttab), 0., ttab) 
     1359    ttab_0 = lbc_add (ttab_0 , nperio=nperio, cd_type='T', psgn=psgn) 
    12201360    iy, ay = __findAxis__ (ttab_0, 'y') 
    12211361    iz, az = __findAxis__ (ttab_0, 'z') 
    1222     vtab = lbc ( 0.5 * (ttab_0 + np.roll (ttab_0, axis=iy, shift=-1)), cd_type='V', nperio=nperio, psgn=psgn) 
    1223  
    1224     if math == xr : 
     1362    if action == 'ave'  : vtab = 0.5 *      (ttab_0 + np.roll (ttab_0, axis=iy, shift=-1)) 
     1363    if action == 'min'  : vtab = np.minimum (ttab_0 , np.roll (ttab_0, axis=iy, shift=-1)) 
     1364    if action == 'max'  : vtab = np.maximum (ttab_0 , np.roll (ttab_0, axis=iy, shift=-1)) 
     1365    if action == 'mult' : vtab =             ttab_0 * np.roll (ttab_0, axis=iy, shift=-1) 
     1366 
     1367    vtab = lbc_del (vtab, nperio=nperio, cd_type='V', psgn=psgn) 
     1368    if mmath == xr : 
    12251369        if ay != None :  
    12261370            vtab = vtab.assign_coords({ay:np.arange(vtab.shape[iy])+1.}) 
    … …  
    12291373    return vtab 
    12301374 
    1231 def V2F (vtab, nperio=None, psgn=-1.0, zdim='depthf') : 
     1375#@numba.jit(forceobj=True) 
     1376def V2F (vtab, nperio=None, psgn=-1.0, zdim='depthf', action='ave') : 
    12321377    '''Interpolate an array from V grid to F grid (i-mean)''' 
    1233     math = __math__ (vtab) 
    1234     vtab_0 = math.where ( np.isnan(vtab), 0., vtab) 
    1235     vtab_0 = lbc (vtab_0 , nperio=nperio, cd_type='V', psgn=psgn) 
     1378    mmath = __mmath__ (vtab) 
     1379    vtab_0 = mmath.where ( np.isnan(vtab), 0., vtab) 
     1380    vtab_0 = lbc_add (vtab_0 , nperio=nperio, cd_type='V', psgn=psgn) 
    12361381    ix, ax = __findAxis__ (vtab_0, 'x') 
    12371382    iz, az = __findAxis__ (vtab_0, 'z') 
    1238     ftab = lbc ( 0.5 * (vtab_0 + np.roll (vtab_0, axis=ix, shift=-1)), cd_type='F', nperio=nperio, psgn=psgn) 
    1239  
    1240     if math == xr : 
     1383    if action == 'ave'  : 0.5 *      (vtab_0 + np.roll (vtab_0, axis=ix, shift=-1)) 
     1384    if action == 'min'  : np.minimum (vtab_0 , np.roll (vtab_0, axis=ix, shift=-1)) 
     1385    if action == 'max'  : np.maximum (vtab_0 , np.roll (vtab_0, axis=ix, shift=-1)) 
     1386    if action == 'mult' :             vtab_0 * np.roll (vtab_0, axis=ix, shift=-1) 
     1387    ftab = lbc_del (ftab, nperio=nperio, cd_type='F', psgn=psgn) 
     1388     
     1389    if mmath == xr : 
    12411390        if ax != None :  
    12421391            ftab = ftab.assign_coords({ax:np.arange(ftab.shape[ix])+1.}) 
    12431392        if zdim != None and iz != None and az != 'olevel' : 
    12441393            ftab = ftab.rename( {az:zdim})  
    1245     return ftab 
    1246  
    1247 def U2F (utab, nperio=None, psgn=-1.0, zdim='depthf') : 
     1394    return lbc_del (ftab, nperio=nperio, cd_type='F', psgn=psgn) 
     1395 
     1396#@numba.jit(forceobj=True) 
     1397def U2F (utab, nperio=None, psgn=-1.0, zdim='depthf', action='ave') : 
    12481398    '''Interpolate an array from U grid to F grid i-mean)''' 
    1249     math = __math__ (utab) 
    1250     utab_0 = math.where ( np.isnan(utab), 0., utab) 
    1251     utab_0 = lbc (utab_0 , nperio=nperio, cd_type='U', psgn=psgn) 
     1399    mmath = __mmath__ (utab) 
     1400    utab_0 = mmath.where ( np.isnan(utab), 0., utab) 
     1401    utab_0 = lbc_add (utab_0 , nperio=nperio, cd_type='U', psgn=psgn) 
    12521402    iy, ay = __findAxis__ (utab_0, 'y') 
    12531403    iz, az = __findAxis__ (utab_0, 'z') 
    1254     ftab = lbc ( 0.5 * (utab_0 + np.roll (utab_0, axis=iy, shift=-1)), cd_type='F', nperio=nperio, psgn=psgn) 
    1255  
    1256     if math == xr : 
     1404    if action == 'ave'  :    ftab = 0.5 *      (utab_0 + np.roll (utab_0, axis=iy, shift=-1)) 
     1405    if action == 'min'  :    ftab = np.minimum (utab_0 , np.roll (utab_0, axis=iy, shift=-1)) 
     1406    if action == 'max'  :    ftab = np.maximum (utab_0 , np.roll (utab_0, axis=iy, shift=-1)) 
     1407    if action == 'mult' :    ftab =             utab_0 * np.roll (utab_0, axis=iy, shift=-1) 
     1408    ftab = lbc_del (ftab, nperio=nperio, cd_type='F', psgn=psgn) 
     1409 
     1410    if mmath == xr : 
    12571411        if ay != None :  
    12581412            ftab = ftab.assign_coords({'y':np.arange(ftab.shape[iy])+1.}) 
    … …  
    12611415    return ftab 
    12621416 
    1263 def F2T (ftab, nperio=None, psgn=1.0, zdim='deptht') : 
     1417#@numba.jit(forceobj=True) 
     1418def F2T (ftab, nperio=None, psgn=1.0, zdim='deptht', action='ave') : 
    12641419    '''Interpolate an array on F grid to T grid (i- and j- means)''' 
    1265     math = __math__ (ftab) 
    1266     ftab_0 = math.where ( np.isnan(ttab), 0., ttab) 
    1267     ftab_0 = lbc (ftab_0 , nperio=nperio, cd_type='F', psgn=psgn) 
    1268     ttab = U2T(F2U(ftab_0, nperio=nperio, psgn=psgn, zdim=zdim), nperio=nperio, psgn=psgn, zdim=zdim) 
    1269     return ttab 
    1270  
    1271 def T2F (ttab, nperio=None, psgn=1.0, zdim='deptht') : 
     1420    mmath = __mmath__ (ftab) 
     1421    ftab_0 = mmath.where ( np.isnan(ttab), 0., ttab) 
     1422    ftab_0 = lbc_add (ftab_0 , nperio=nperio, cd_type='F', psgn=psgn) 
     1423    ttab = U2T(F2U(ftab_0, nperio=nperio, psgn=psgn, zdim=zdim, action=action), nperio=nperio, psgn=psgn, zdim=zdim, action=action) 
     1424    return lbc_del (ttab, nperio=nperio, cd_type='T', psgn=psgn) 
     1425 
     1426#@numba.jit(forceobj=True) 
     1427def T2F (ttab, nperio=None, psgn=1.0, zdim='deptht', action='mean') : 
    12721428    '''Interpolate an array on T grid to F grid (i- and j- means)''' 
    1273     math = __math__ (ftab) 
    1274     ttab_0 = math.where ( np.isnan(ttab), 0., ttab) 
    1275     ttab_0 = lbc (ttab_0 , nperio=nperio, cd_type='T', psgn=psgn) 
    1276     ftab = T2U(U2F(ttab, nperio=nperio, psgn=psgn, zdim=zdim), nperio=nperio, psgn=psgn, zdim=zdim) 
    1277     return ftab 
    1278  
    1279 def F2U (ftab, nperio=None, psgn=1.0, zdim='depthu') : 
     1429    mmath = __mmath__ (ttab) 
     1430    ttab_0 = mmath.where ( np.isnan(ttab), 0., ttab) 
     1431    ttab_0 = lbc_add (ttab_0 , nperio=nperio, cd_type='T', psgn=psgn) 
     1432    ftab = T2U(U2F(ttab, nperio=nperio, psgn=psgn, zdim=zdim, action=action), nperio=nperio, psgn=psgn, zdim=zdim, action=action) 
     1433     
     1434    return lbc_del (ftab, nperio=nperio, cd_type='F', psgn=psgn) 
     1435 
     1436#@numba.jit(forceobj=True) 
     1437def F2U (ftab, nperio=None, psgn=1.0, zdim='depthu', action='ave') : 
    12801438    '''Interpolate an array on F grid to FUgrid (i-mean)''' 
    1281     math = __math__ (ftab) 
    1282     ftab_0 = math.where ( np.isnan(ftab), 0., ftab) 
    1283     ftab_0 = lbc (ftab_0 , nperio=nperio, cd_type='F', psgn=psgn) 
     1439    mmath = __mmath__ (ftab) 
     1440    ftab_0 = mmath.where ( np.isnan(ftab), 0., ftab) 
     1441    ftab_0 = lbc_add (ftab_0 , nperio=nperio, cd_type='F', psgn=psgn) 
    12841442    iy, ay = __findAxis__ (ftab_0, 'y') 
    12851443    iz, az = __findAxis__ (ftab_0, 'z') 
    1286     utab = lbc ( 0.5 * (ftab_0 + np.roll (ftab_0, axis=iy, shift=-1)), cd_type='U', nperio=nperio, psgn=psgn) 
    1287      
    1288     if math == xr : 
     1444    if action == 'ave'  : utab = 0.5 *      (ftab_0 + np.roll (ftab_0, axis=iy, shift=-1)) 
     1445    if action == 'min'  : utab = np.minimum (ftab_0 , np.roll (ftab_0, axis=iy, shift=-1)) 
     1446    if action == 'max'  : utab = np.maximum (ftab_0 , np.roll (ftab_0, axis=iy, shift=-1)) 
     1447    if action == 'mult' : utab =             ftab_0 * np.roll (ftab_0, axis=iy, shift=-1) 
     1448 
     1449    utab = lbc_del (utab, nperio=nperio, cd_type='U', psgn=psgn) 
     1450     
     1451    if mmath == xr : 
    12891452        utab = utab.assign_coords({ay:np.arange(ftab.shape[iy])+1.}) 
    12901453        if zdim != None and iz != None and az != 'olevel' : 
    … …  
    12921455    return utab 
    12931456 
    1294 def F2V (ftab, nperio=None, psgn=1.0, zdim='depthv') : 
     1457#@numba.jit(forceobj=True) 
     1458def F2V (ftab, nperio=None, psgn=1.0, zdim='depthv', action='ave') : 
    12951459    '''Interpolate an array from F grid to V grid (i-mean)''' 
    1296     math = __math__ (ftab) 
    1297     ftab_0 = math.where ( np.isnan(ftab), 0., ftab) 
    1298     ftab_0 = lbc (ftab_0 , nperio=nperio, cd_type='F', psgn=psgn) 
     1460    mmath = __mmath__ (ftab) 
     1461    ftab_0 = mmath.where ( np.isnan(ftab), 0., ftab) 
     1462    ftab_0 = lbc_add (ftab_0 , nperio=nperio, cd_type='F', psgn=psgn) 
    12991463    ix, ax = __findAxis__ (ftab_0, 'x') 
    13001464    iz, az = __findAxis__ (ftab_0, 'z') 
    1301     vtab = lbc ( 0.5 * (ftab_0 + np.roll (ftab_0, axis=ix, shift=-1)), cd_type='V', nperio=nperio) 
    1302  
    1303     if math == xr : 
     1465    if action == 'ave'  : vtab = 0.5 *      (ftab_0 + np.roll (ftab_0, axis=ix, shift=-1)) 
     1466    if action == 'min'  : vtab = np.minimum (ftab_0 , np.roll (ftab_0, axis=ix, shift=-1)) 
     1467    if action == 'max'  : vtab = np.maximum (ftab_0 , np.roll (ftab_0, axis=ix, shift=-1)) 
     1468    if action == 'mult' : vtab =             ftab_0 * np.roll (ftab_0, axis=ix, shift=-1) 
     1469 
     1470    vtab = lbc_del (vtab, nperio=nperio, cd_type='V', psgn=psgn) 
     1471    if mmath == xr : 
    13041472        vtab = vtab.assign_coords({ax:np.arange(ftab.shape[ix])+1.}) 
    13051473        if zdim != None and iz != None and az != 'olevel' : 
    … …  
    13071475    return vtab 
    13081476 
     1477#@numba.jit(forceobj=True) 
    13091478def W2T (wtab, zcoord=None, zdim='deptht', sval=np.nan) : 
    13101479    ''' 
    … …  
    13121481    sval is the bottom value 
    13131482    ''' 
    1314     math = __math__ (wtab) 
    1315     wtab_0 = math.where ( np.isnan(wtab), 0., wtab) 
     1483    mmath = __mmath__ (wtab) 
     1484    wtab_0 = mmath.where ( np.isnan(wtab), 0., wtab) 
    13161485 
    13171486    iz, az = __findAxis__ (wtab_0, 'z') 
    … …  
    13191488    ttab = 0.5 * ( wtab_0 + np.roll (wtab_0, axis=iz, shift=-1) ) 
    13201489     
    1321     if math == xr : 
     1490    if mmath == xr : 
    13221491        ttab[{az:iz}] = sval 
    13231492        if zdim != None and iz != None and az != 'olevel' : 
    … …  
    13301499    return ttab 
    13311500 
     1501#@numba.jit(forceobj=True) 
    13321502def T2W (ttab, zcoord=None, zdim='depthw', sval=np.nan, extrap_surf=False) : 
    13331503    '''Interpolate an array from T grid to W grid (k-mean) 
    … …  
    13351505    if extrap_surf==True, surface value is taken from 1st level value. 
    13361506    ''' 
    1337     math = __math__ (ttab) 
    1338     ttab_0 = math.where ( np.isnan(ttab), 0., ttab) 
     1507    mmath = __mmath__ (ttab) 
     1508    ttab_0 = mmath.where ( np.isnan(ttab), 0., ttab) 
    13391509    iz, az = __findAxis__ (ttab_0, 'z') 
    13401510    wtab = 0.5 * ( ttab_0 + np.roll (ttab_0, axis=iz, shift=1) ) 
    13411511 
    1342     if math == xr : 
     1512    if mmath == xr : 
    13431513        if extrap_surf : wtab[{az:0}] = ttabb[{az:0}] 
    13441514        else           : wtab[{az:0}] = sval 
    … …  
    13471517        else           : wtab[..., 0, :, :] = sval 
    13481518 
    1349     if math == xr : 
     1519    if mmath == xr : 
    13501520        if zdim != None and iz != None and az != 'olevel' : 
    13511521                wtab = wtab.rename ( {az:zdim}) 
    … …  
    13541524    return wtab 
    13551525 
    1356 def fill (ptab, nperio, cd_type='T', npass=1) : 
    1357     ''' 
    1358     Fill 0. or np.nan values with mean of neighbours 
     1526#@numba.jit(forceobj=True) 
     1527def fill (ptab, nperio, cd_type='T', npass=1, sval=0.) : 
     1528    ''' 
     1529    Fill sval values with mean of neighbours 
    13591530    
    13601531    Inputs : 
    … …  
    13631534    '''        
    13641535 
    1365     math = __math__ (ptab) 
    1366     ztab   = math.where (np.isnan(ptab), 0., ptab) 
    1367  
    1368     DoPerio = False 
    1369     if nperio == 4.2 or nperio == 6.2 : DoPerio = True 
    1370          
    1371     if DoPerio :  ztab = lbc_add (ztab) 
     1536    mmath = __mmath__ (ptab) 
     1537 
     1538    DoPerio = False ; lperio = nperio 
     1539    if nperio == 4.2 : 
     1540        DoPerio = True ; lperio = 4 
     1541    if nperio == 6.2 : 
     1542        DoPerio = True ; lperio = 6 
     1543         
     1544    if DoPerio : 
     1545        ztab = lbc_add (ptab, nperio=nperio, sval=sval) 
     1546    else :  
     1547        ztab = ptab 
     1548         
     1549    if np.isnan (sval) :  
     1550        ztab   = mmath.where (np.isnan(ztab), np.nan, ztab) 
     1551    else : 
     1552        ztab   = mmath.where (ztab==sval    , np.nan, ztab) 
    13721553    
    13731554    for nn in np.arange (npass) :  
    1374         zmask  = math.where (ztab==0., 0., 1.  ) 
     1555        zmask = mmath.where ( np.isnan(ztab), 0., 1.   ) 
     1556        ztab0 = mmath.where ( np.isnan(ztab), 0., ztab ) 
    13751557        # Compte du nombre de voisins 
    1376         zcount = zmask \ 
     1558        zcount = 1./6. * ( zmask \ 
    13771559          + np.roll(zmask, shift=1, axis=-1) + np.roll(zmask, shift=-1, axis=-1) \ 
    13781560          + np.roll(zmask, shift=1, axis=-2) + np.roll(zmask, shift=-1, axis=-2) \ 
    … …  
    13811563                + np.roll(np.roll(zmask, shift=-1, axis=-2), shift= 1, axis=-1) \ 
    13821564                + np.roll(np.roll(zmask, shift= 1, axis=-2), shift=-1, axis=-1) \ 
    1383                 + np.roll(np.roll(zmask, shift=-1, axis=-2), shift=-1, axis=-1) ) 
    1384  
    1385         znew = zmask \ 
    1386           + np.roll(ztab, shift=1, axis=-1) + np.roll(ztab, shift=-1, axis=-1) \ 
    1387           + np.roll(ztab, shift=1, axis=-2) + np.roll(ztab, shift=-1, axis=-2) \ 
    1388           + 0.5 * ( \ 
    1389             + np.roll(np.roll(ztab, shift= 1, axis=-2), shift= 1, axis=-1) \ 
    1390             + np.roll(np.roll(ztab, shift=-1, axis=-2), shift= 1, axis=-1) \ 
    1391             + np.roll(np.roll(ztab, shift= 1, axis=-2), shift=-1, axis=-1) \ 
    1392             + np.roll(np.roll(ztab, shift=-1, axis=-2), shift=-1, axis=-1) ) 
    1393  
    1394         zcount = lbc (zcount, nperio=nperio, cd_type=cd_type) 
    1395         znew   = lbc (znew  , nperio=nperio, cd_type=cd_type) 
    1396      
    1397         ztab = math.where (np.logical_and (zmask==0., zcount>0), znew/zcount, ztab) 
    1398          
    1399     ztab = math.where (ztab==0., np.nan, ztab) 
    1400  
    1401     if DoPerio : ztab = lbc_del (ztab) 
     1565                + np.roll(np.roll(zmask, shift=-1, axis=-2), shift=-1, axis=-1) ) ) 
     1566 
     1567        znew =1./6. * ( ztab0 \ 
     1568           + np.roll(ztab0, shift=1, axis=-1) + np.roll(ztab0, shift=-1, axis=-1) \ 
     1569           + np.roll(ztab0, shift=1, axis=-2) + np.roll(ztab0, shift=-1, axis=-2) \ 
     1570           + 0.5 * ( \ 
     1571                + np.roll(np.roll(ztab0 , shift= 1, axis=-2), shift= 1, axis=-1) \ 
     1572                + np.roll(np.roll(ztab0 , shift=-1, axis=-2), shift= 1, axis=-1) \ 
     1573                + np.roll(np.roll(ztab0 , shift= 1, axis=-2), shift=-1, axis=-1) \ 
     1574                + np.roll(np.roll(ztab0 , shift=-1, axis=-2), shift=-1, axis=-1) ) ) 
     1575 
     1576        zcount = lbc (zcount, nperio=lperio, cd_type=cd_type) 
     1577        znew   = lbc (znew  , nperio=lperio, cd_type=cd_type) 
     1578         
     1579        ztab = mmath.where (np.logical_and (zmask==0., zcount>0), znew/zcount, ztab) 
     1580 
     1581    ztab = mmath.where (zcount==0, sval, ztab) 
     1582    if DoPerio : ztab = lbc_del (ztab, nperio=lperio) 
    14021583 
    14031584    return ztab 
    14041585 
     1586#@numba.jit(forceobj=True) 
    14051587def correct_uv (u, v, lat) : 
    14061588    ''' 
    … …  
    14171599 
    14181600    Outputs :  
    1419        modified eastward/nothward components have correct polar projection in cartopy 
    1420     ''' 
    1421     math = __math__ (u) 
     1601       modified eastward/nothward components to have correct polar projections in cartopy 
     1602    ''' 
    14221603    uv = np.sqrt (u*u + v*v)           # Original modulus 
    14231604    zu = u 
    … …  
    14271608    return uc, vc 
    14281609 
     1610def msf (v_e1v_e3v, lat1d, depthw) : 
     1611    ''' 
     1612    Computes the meridonal stream function 
     1613    First input is meridional_velocity*e1v*e3v 
     1614    ''' 
     1615    @numba.jit(forceobj=True) 
     1616    def iin (tab, dim) : 
     1617        ''' 
     1618        Integrate from the bottom 
     1619        ''' 
     1620        result = tab * 0.0 
     1621        nlen = len(tab.coords[dim]) 
     1622        for jn in np.arange (nlen-2, 0, -1) : 
     1623            result [{dim:jn}] = result [{dim:jn+1}] - tab [{dim:jn}] 
     1624        result = result.where (result !=0, np.nan) 
     1625        return result 
     1626     
     1627    zomsf = iin ((v_e1v_e3v).sum (dim='x', keep_attrs=True)*1E-6, dim='depthv') 
     1628    zomsf = zomsf.assign_coords ( {'depthv':depthw.values, 'y':lat1d}) 
     1629    zomsf = zomsf.rename ( {'depthv':'depthw', 'y':'lat'}) 
     1630    zomsf.attrs['long_name'] = 'Meridional stream function' 
     1631 
     1632    zomsf.attrs['units'] = 'Sv' 
     1633    zomsf.depthw.attrs=depthw.attrs 
     1634    zomsf.lat.attrs=lat1d.attrs 
     1635         
     1636    return zomsf 
     1637 
     1638def bsf (u_e2u_e3u, mask, nperio=None, bsf0=None ) : 
     1639    ''' 
     1640    Computes the barotropic stream function 
     1641    First input is zonal_velocity*e2u*e3u 
     1642    bsf0 is the point with bsf=0 (ex: bsf0={'x':5, 'y':120} ) 
     1643    ''' 
     1644    @numba.jit(forceobj=True) 
     1645    def iin (tab, dim) : 
     1646        ''' 
     1647        Integrate from the south 
     1648        ''' 
     1649        result = tab * 0.0 
     1650        nlen = len(tab.coords[dim]) 
     1651        for jn in np.arange (3, nlen) : 
     1652            result [{dim:jn}] = result [{dim:jn-1}] + tab [{dim:jn}] 
     1653        return result 
     1654     
     1655    bsf = iin ((u_e2u_e3u).sum(dim='depthu', keep_attrs=True)*1E-6, dim='y') 
     1656    bsf.attrs = u_e2u_e3u.attrs 
     1657    if bsf0 != None : 
     1658        bsf = bsf - bsf.isel (bsf0) 
     1659        
     1660    bsf = bsf.where (mask !=0, np.nan) 
     1661    bsf.attrs['long_name'] = 'Barotropic stream function' 
     1662    bsf.attrs['units'] = 'Sv' 
     1663    bsf = lbc (bsf, nperio=nperio, cd_type='F') 
     1664        
     1665    return bsf 
     1666 
     1667def namelist_read (ref=None, cfg=None, out='dict', flat=False, verbose=False) : 
     1668    ''' 
     1669    Read NEMO namelist(s) and return either a dictionnary or an xarray dataset 
     1670 
     1671    ref : file with reference namelist, or a f90nml.namelist.Namelist object 
     1672    cfg : file with config namelist, or a f90nml.namelist.Namelist object 
     1673    At least one namelist neaded 
     1674 
     1675    out:  
     1676        'dict' to return a dictonnary 
     1677        'xr'   to return an xarray dataset 
     1678    flat : only for dict output. Output a flat dictionnary with all values. 
     1679     
     1680    ''' 
     1681 
     1682    if ref != None : 
     1683        if isinstance (ref, str) : nml_ref = f90nml.read (ref) 
     1684        if isinstance (ref, f90nml.namelist.Namelist) : nml_ref = ref 
     1685         
     1686    if cfg != None : 
     1687        if isinstance (cfg, str) : nml_cfg = f90nml.read (cfg) 
     1688        if isinstance (cfg, f90nml.namelist.Namelist) : nml_cfg = cfg 
     1689     
     1690    if out == 'dict' : dict_namelist = {} 
     1691    if out == 'xr'   : xr_namelist = xr.Dataset () 
     1692 
     1693    list_nml = [] ; list_comment = [] 
     1694 
     1695    if ref != None : 
     1696        list_nml.append (nml_ref) ; list_comment.append ('ref') 
     1697    if cfg != None : 
     1698        list_nml.append (nml_cfg) ; list_comment.append ('cfg') 
     1699 
     1700    for nml, comment in zip (list_nml, list_comment) : 
     1701        if verbose : print (comment) 
     1702        if flat and out =='dict' : 
     1703            for nam in nml.keys () : 
     1704                if verbose : print (nam) 
     1705                for value in nml[nam] : 
     1706                     if out == 'dict' : dict_namelist[value] = nml[nam][value] 
     1707                     if verbose : print (nam, ':', value, ':', nml[nam][value]) 
     1708        else : 
     1709            for nam in nml.keys () : 
     1710                if verbose : print (nam) 
     1711                if out == 'dict' : 
     1712                    if nam not in dict_namelist.keys () : dict_namelist[nam] = {} 
     1713                for value in nml[nam] : 
     1714                    if out == 'dict' : dict_namelist[nam][value] = nml[nam][value] 
     1715                    if out == 'xr'   : xr_namelist[value] = nml[nam][value] 
     1716                    if verbose : print (nam, ':', value, ':', nml[nam][value]) 
     1717 
     1718    if out == 'dict' : return dict_namelist 
     1719    if out == 'xr'   : return xr_namelist 
     1720 
     1721 
     1722def fill_closed_seas (imask, nperio=None,  cd_type='T') : 
     1723    '''Fill closed seas with image processing library 
     1724    imask : mask, 1 on ocean, 0 on land 
     1725    ''' 
     1726    from scipy import ndimage 
     1727 
     1728    imask_filled = ndimage.binary_fill_holes ( lbc (imask, nperio=nperio, cd_type=cd_type)) 
     1729    imask_filled = lbc ( imask_filled, nperio=nperio, cd_type=cd_type) 
     1730 
     1731    return imask_filled 
     1732 
    14291733## =========================================================================== 
    14301734## 
    … …  
    14331737## =========================================================================== 
    14341738 
    1435 def is_orca_north_fold ( Xtest, cname_long='T' ) : 
     1739def __is_orca_north_fold__ ( Xtest, cname_long='T' ) : 
    14361740    ''' 
    14371741    Ported (pirated !!?) from Sosie 
Note: See TracChangeset for help on using the changeset viewer.