Changeset 5948


Ignore:
Timestamp:
10/14/21 16:14:50 (3 years ago)
Author:
snguyen
Message:

uptade to new nemo.py module with correction of a bug in U direction for periodicity 4 in LBC. Minor changes in Build_coordinates_mask.py script.

Location:
TOOLS/MOSAIX
Files:
2 edited

Legend:

Unmodified
Added
Removed

  • TOOLS/MOSAIX/Build_coordinates_mask.py

    r5947 r5948  
    3232parser.add_argument ('--coord'   , help='coordinates file name', type=str, default='coordinates.nc' ) 
    3333parser.add_argument ('--fout'    , help='output file name (given as an input to MOSAIX)', type=str, default='coordinates_mask.nc' ) 
    34 parser.add_argument ('--ocePerio', help='periodicity of ocean grid', type=int, default=0, choices=(1, 2, 3, 4, 5, 6) ) 
     34parser.add_argument ('--ocePerio', help='periodicity of ocean grid', type=int, default=0, choices=range(1, 7) ) 
    3535parser.add_argument ('--nemo4Ubug',help='reproduce NEMO lbc_lnk bug for U grid and periodicity 4', type=bool, default=False, choices=(True, False) ) 
    3636parser.add_argument ('--straits' , help='modify grid metrics for selected strait points (depends on model name)', type=bool, default=False, choices=(True, False) ) 

  • TOOLS/MOSAIX/nemo.py

    r4259 r5948  
    1 ### =========================================================================== 
    2 ### 
    3 ### Periodicity of ORCA fields 
    4 ### 
    5 ### =========================================================================== 
     1# -*- coding: utf-8 -*- 
     2## =========================================================================== 
    63## 
    74##  Warning, to install, configure, run, use any of Olivier Marti's 
    … …  
    1411##  personal. 
    1512## 
     13''' 
     14Utilities to plot NEMO ORCA fields 
     15Periodicity and other stuff 
     16 
     17olivier.marti@lsce.ipsl.fr 
     18''' 
     19 
     20import sys 
     21 
     22try    : import numpy as np 
     23except : pass 
     24 
     25try    : import xarray as xr 
     26except : pass 
     27 
    1628## SVN information 
    1729__Author__   = "$Author$" 
    … …  
    2133__HeadURL    = "$HeadURL$" 
    2234 
    23 import sys, numpy as np 
    24  
    25 def __guessNperio__ (jpi, nperio) : 
    26     """ 
    27     Tries to guess the value of nperio 
    28     """ 
     35def __guessNperio__ (jpi, nperio=None) : 
     36    ''' 
     37    Tries to guess the value of nperio (periodicity parameter. See NEMO documentation for details) 
     38     
     39    Inputs 
     40    jpi    : number of longitudes 
     41    nperio : periodicity parameter 
     42    ''' 
     43    
    2944    if nperio == None : 
    30         if jpi ==  182 : nperio = 4 #   ORCA2. We choose legacy orca2 
     45        if jpi ==  182 : nperio = 4 #   ORCA2. We choose legacy orca2. 
    3146        if jpi ==  362 : nperio = 6 #   ORCA1. 
    3247        if jpi == 1442 : nperio = 6 # ORCA025. 
    … …  
    3550            sys.exit ('in nemo.lbc : nperio not found, and cannot by guessed' ) 
    3651        else : 
    37             print ('nperio set as {:d} (deduced from jpi={:d}'.format(nperio, jpi)) 
    38  
    39     return nperio     
    40              
    41 def lbc (ptab, nperio=None, cd_type='T', psgn=1.0) : 
     52            print ('nperio set as {:d} (deduced from jpi={:d})'.format(nperio, jpi)) 
     53 
     54    return nperio 
     55 
     56def fixed_lon (lon) : 
     57    ''' 
     58    Returns corrected longitudes for nicer plots 
     59 
     60    fixed_lon (lon) 
     61    lon : longitudes of the grid. At least 2D. 
     62    ''' 
     63    fixed_lon = lon.copy () 
     64    for i, start in enumerate (np.argmax (np.abs (np.diff (lon, axis=-1)) > 180, axis=-1)) : 
     65        fixed_lon[...,i, start+1:] += 360 
     66    return fixed_lon 
     67 
     68def jeq (lat) : 
     69    ''' 
     70    Returns j index of equator in the grid 
     71     
     72    lat : latitudes of the grid. At least 2D. 
     73    ''' 
     74    jeq = np.argmin (np.abs (np.float64(lat[:,0]))) 
     75    return jeq 
     76 
     77def lon1D (lon, lat=None) : 
     78    ''' 
     79    Returns 1D longitude for simple plots. 
     80     
     81    lon : longitudes of the grid 
     82    lat (optionnal) : latitudes of the grid 
     83    ''' 
     84     
     85    if np.max (lat) != None : 
     86        je     = jeq (lat) 
     87        lon1D = lon[..., je, :] 
     88    else : 
     89        jpj, jpi = lon.shape[-2:] 
     90        lon1D = lon[..., jpj//3, :] 
     91 
     92    start = np.argmax (np.abs (np.diff (lon1D, axis=-1)) > 180.0, axis=-1) 
     93    lon1D[..., start+1:] += 360 
     94         
     95    return lon1D 
     96 
     97def latreg (lat) : 
     98    ''' 
     99    Returns maximum j index wehre gridlines are along latitude in the norethern hemisphere 
     100     
     101    lon : longitudes of the grid 
     102    lat (optionnal) : latitudes of the grid 
     103    ''' 
     104    dy     = np.float64 (np.mean (np.abs (lat - np.roll(lat,shift=1,axis=-2)))) 
     105    je     = jeq (lat) 
     106    jreg   = np.where (lat[...,je:,:].max(axis=-1) - lat[...,je:,:].min(axis=-1)< dy/100.)[-1][-1] + je 
     107    latreg = np.float64 (lat[...,jreg,:].mean(axis=-1)) 
     108    JREG = jreg 
     109 
     110    return jreg, latreg 
     111 
     112def lat1D (lat) : 
     113    ''' 
     114    Returns 1D latitudes for zonal means and simple plots. 
     115 
     116    lat : latitudes of the grid 
     117    ''' 
     118    jpj, jpi = lat.shape[-2:] 
     119    try : 
     120        if type (lat) ==  xr.core.dataarray.DataArray : math = xr 
     121        else : math = np 
     122    except : math = np 
     123 
     124    # jreg : index of last uniform latitude, north of equator 
     125 
     126    dy = np.float64 (np.mean (np.abs (lat - np.roll (lat, shift=1,axis=-2)))) 
     127    je     = jeq (lat) 
     128    lateq =  np.float64 (lat[...,je,:].mean(axis=-1)) 
     129       
     130    jreg   = np.where (lat[...,je:,:].max(axis=-1) - lat[...,je:,:].min(axis=-1)< dy/100.)[-1][-1] + je 
     131    latreg = np.float64 (lat[...,jreg,:].mean(axis=-1)) 
     132    latave = np.mean (lat, axis=-1) 
     133 
     134    if ( np.abs (lateq) <  dy/100. ) : # T, U or W grid 
     135        dys = (90.-latreg)//(jpj-jreg-1)*0.5 
     136        yrange = 90.-dys-latreg 
     137    else                          :  # V or F grid 
     138        yrange = 90.    -latreg 
     139         
     140    lat1D = math.where (latave<latreg, latave, latreg + yrange * (np.arange(jpj)-jreg)/(jpj-jreg-1) )    
     141         
     142    if math == xr : 
     143        lat1D.attrs = lat.attrs 
     144 
     145    return lat1D 
     146 
     147def latlon1D (lat, lon) : 
     148    ''' 
     149    Returns simple latitude and longitude (1D) for simple plots. 
     150    ''' 
     151    laty = lat1D (lat) 
     152    lonx = lon1D (lon, lat) 
     153 
     154    return laty, lonx 
     155         
     156def extend (tab, Lon=False, jplus=25) : 
     157    ''' 
     158    Returns extended field eastward to have better plots 
     159    Works only for xarray and numpy data (?) 
     160 
     161    tab : field to extend. 
     162    Lon : (optional, default=False) : if True, add 360 in the extended parts of the field 
     163    jplus (optional, default=25) : number of points added on the east side of the field 
     164     
     165    ''' 
     166    jpi   = tab.shape[-1] 
     167    JPLUS = jplus 
     168     
     169    try :  
     170        if type (tab) ==  xr.core.dataarray.DataArray : math = xr 
     171        else : math = np 
     172    except : math = np 
     173 
     174    if Lon : xplus = -360.0 
     175    else   : xplus =    0.0 
     176 
     177    if math == xr :  
     178        extend = xr.concat      ( (tab[..., 2:jpi] + xplus, tab[..., 2:jplus]), dim=tab.dims[-1] ) 
     179    if math == np : 
     180        extend = np.concatenate ( (tab[..., 2:jpi] + xplus, tab[..., 2:jplus]), axis=-1 ) 
     181             
     182    return extend 
     183 
     184         
     185def lbc (ptab, nperio=None, cd_type='T', psgn=1.0, nemo_4U_bug=False) : 
    42186    """ 
    43     Set periodicity on input fields 
    44     ptab      : Input array 
     187    Set periodicity on input field 
     188    ptab      : Input array (works  
    45189      rank 2 at least : ptab[...., lat, lon] 
    46190    nperio    : Type of periodicity 
    … …  
    54198    See NEMO documentation for further details 
    55199    """ 
    56          
     200 
    57201    jpi    = ptab.shape[-1] 
    58     nperio = __guessNperio__ ( jpi, nperio ) 
     202    nperio = __guessNperio__ ( jpi, nperio) 
    59203    psgn   = ptab.dtype.type(psgn) 
    60204     
    … …  
    62206    #> East-West boundary conditions 
    63207    # ------------------------------ 
    64  
    65208    if nperio in [1, 4, 6] : 
    66209        # ... cyclic 
    67         ptab [...,:,  0] = ptab [...,:,-2] 
    68         ptab [...,:, -1] = ptab [...,:, 1] 
     210        ptab [..., :,  0] = ptab [..., :, -2] 
     211        ptab [..., :, -1] = ptab [..., :, 1] 
    69212    
    70213    # 
    … …  
    73216    if nperio in [3, 4] :  # North fold T-point pivot 
    74217        if cd_type in [ 'T', 'W' ] : # T-, W-point 
    75             ptab[..., -1, 1:       ] = psgn * ptab[..., -3, -1:0:-1      ]       
    76             ptab[..., -1, 0        ] = psgn * ptab[..., -3, 2            ] 
    77             ptab[..., -2, jpi//2:  ] = psgn * ptab[..., -2, jpi//2:0:-1  ] 
     218            ptab [..., -1, 1:       ] = psgn * ptab [..., -3, -1:0:-1      ]       
     219            ptab [..., -1, 0        ] = psgn * ptab [..., -3, 2            ] 
     220            ptab [..., -2, jpi//2:  ] = psgn * ptab [..., -2, jpi//2:0:-1  ] 
    78221                       
    79222        if cd_type == 'U' : 
    80             ptab[..., -1, 0:-1     ] = psgn * ptab[..., -3, -1:0:-1      ]        
    81             ptab[..., -1,  0       ] = psgn * ptab[..., -3,  1           ] 
    82             ptab[..., -1, -1       ] = psgn * ptab[..., -3, -2           ]  
    83             ptab[..., -2, jpi//2-1:] = psgn * ptab[..., -2, jpi//2+1:0:-1] 
     223            ptab [..., -1, 0:-1     ] = psgn * ptab [..., -3, -1:0:-1      ]        
     224            ptab [..., -1,  0       ] = psgn * ptab [..., -3,  1           ] 
     225            ptab [..., -1, -1       ] = psgn * ptab [..., -3, -2           ] 
     226 
     227            if nemo_4U_bug : 
     228                ptab [..., -2, jpi//2+1:-1] = psgn * ptab [..., -2, jpi//2-2:0:-1] 
     229                ptab [..., -2, jpi//2-1   ] = psgn * ptab [..., -2, jpi//2     ] 
     230            else : 
     231                ptab [..., -2, jpi//2-1:-1] = psgn * ptab [..., -2, jpi//2:0:-1] 
    84232             
    85233        if cd_type == 'V' :  
    86             ptab[..., -2, 1:       ] = psgn * ptab[..., -3, jpi-1:0:-1   ] 
    87             ptab[..., -1, 1:       ] = psgn * ptab[..., -4, -1:0:-1      ]    
    88             ptab[..., -1, 0        ] = psgn * ptab[..., -4, 2            ] 
     234            ptab [..., -2, 1:       ] = psgn * ptab [..., -3, jpi-1:0:-1   ] 
     235            ptab [..., -1, 1:       ] = psgn * ptab [..., -4, -1:0:-1      ]    
     236            ptab [..., -1, 0        ] = psgn * ptab [..., -4, 2            ] 
    89237             
    90238        if cd_type == 'F' : 
    91             ptab[..., -2, 0:-1     ] = psgn * ptab[..., -3, -1:0:-1      ] 
    92             ptab[..., -1, 0:-1     ] = psgn * ptab[..., -4, -1:0:-1      ] 
    93             ptab[..., -1,  0       ] = psgn * ptab[..., -4,  1           ] 
    94             ptab[..., -1, -1       ] = psgn * ptab[..., -4, -2           ]  
     239            ptab [..., -2, 0:-1     ] = psgn * ptab [..., -3, -1:0:-1      ] 
     240            ptab [..., -1, 0:-1     ] = psgn * ptab [..., -4, -1:0:-1      ] 
     241            ptab [..., -1,  0       ] = psgn * ptab [..., -4,  1           ] 
     242            ptab [..., -1, -1       ] = psgn * ptab [..., -4, -2           ]  
    95243       
    96244    if nperio in [5, 6] :            #  North fold F-point pivot   
    97245        if cd_type in ['T', 'W']  : 
    98             ptab[..., -1, 0:       ] = psgn * ptab[..., -2, -1::-1       ] 
     246            ptab [..., -1, 0:       ] = psgn * ptab [..., -2, -1::-1       ] 
    99247                 
    100248        if cd_type == 'U' : 
    101             ptab[..., -1, 0:-1     ] = psgn * ptab[..., -2, -2::-1       ]        
    102             ptab[..., -1, -1       ] = psgn * ptab[..., -2, 0            ] 
     249            ptab [..., -1, 0:-1     ] = psgn * ptab [..., -2, -2::-1       ]        
     250            ptab [..., -1, -1       ] = psgn * ptab [..., -2, 0            ] 
    103251              
    104252        if cd_type == 'V' : 
    105             ptab[..., -1, 0:       ] = psgn * ptab[..., -3, -1::-1       ] 
    106             ptab[..., -2, jpi//2:  ] = psgn * ptab[..., -2, jpi//2-1::-1 ] 
     253            ptab [..., -1, 0:       ] = psgn * ptab [..., -3, -1::-1       ] 
     254            ptab [..., -2, jpi//2:  ] = psgn * ptab [..., -2, jpi//2-1::-1 ] 
    107255                              
    108256        if cd_type == 'F' : 
    109             ptab[..., -1, 0:-1     ] = psgn * ptab[..., -3, -2::-1       ] 
    110             ptab[..., -1, -1       ] = psgn * ptab[..., -3, 0            ] 
    111             ptab[..., -2, jpi//2:-1] = psgn * ptab[..., -2, jpi//2-2::-1 ] 
     257            ptab [..., -1, 0:-1     ] = psgn * ptab [..., -3, -2::-1       ] 
     258            ptab [..., -1, -1       ] = psgn * ptab [..., -3, 0            ] 
     259            ptab [..., -2, jpi//2:-1] = psgn * ptab [..., -2, jpi//2-2::-1 ] 
    112260 
    113261    return ptab 
    114262 
    115  
    116 def lbc_mask (ptab, nperio=None, cd_type='T') : 
     263def lbc_mask (ptab, nperio=None, cd_type='T', sval=None) : 
    117264    # 
    118265    """ 
    … …  
    131278 
    132279    jpi    = ptab.shape[-1] 
    133     nperio = __guessNperio__ ( jpi, nperio ) 
    134     zero   = ptab.dtype.type(0) 
    135      
     280    nperio = __guessNperio__ (jpi, nperio) 
     281    if sval == None : sval = ptab.dtype.type (0) 
     282 
    136283    # 
    137284    #> East-West boundary conditions 
    … …  
    139286    if nperio in [1, 4, 6] : 
    140287        # ... cyclic 
    141         ptab [...,:,  0] = zero 
    142         ptab [...,:, -1] = zero 
     288        ptab [...,:,  0] = sval 
     289        ptab [...,:, -1] = sval 
    143290 
    144291    # 
    … …  
    146293    # -------------------------------- 
    147294    if nperio in [1, 3, 4, 5, 6] : 
    148         ptab[...,0,:] = zero 
     295        ptab [...,0,:] = sval 
    149296         
    150297    # 
    151298    #> North-South boundary conditions 
    152299    # -------------------------------- 
    153     if nperio in [3, 4] :  # North fold T-point pivot      
     300    if nperio in [3, 4] :  # North fold T-point pivot 
    154301        if cd_type in [ 'T', 'W' ] : # T-, W-point 
    155             ptab[..., -1, 1:       ] = zero 
    156             ptab[..., -1, 0        ] = zero  
    157             ptab[..., -2, jpi//2:  ] = zero 
     302            ptab [..., -1, 1:       ] = sval 
     303            ptab [..., -1, 0        ] = sval  
     304            ptab [..., -2, jpi//2:  ] = sval 
    158305                       
    159306        if cd_type == 'U' : 
    160             ptab[..., -1, 0:-1     ] = zero   
    161             ptab[..., -1,  0       ] = zero 
    162             ptab[..., -1, -1       ] = zero  
    163             ptab[..., -2, jpi//2-1:] = zero 
    164              
    165         if cd_type == 'V' :  
    166             ptab[..., -2, 1:       ] = zero 
    167             ptab[..., -1, 1:       ] = zero    
    168             ptab[..., -1, 0        ] = zero 
     307            ptab [..., -1, 0:-1     ] = sval   
     308            ptab [..., -1,  0       ] = sval 
     309            ptab [..., -1, -1       ] = sval  
     310            ptab [..., -2, jpi//2  :] = sval 
     311             
     312        if cd_type == 'V' : 
     313            ptab [..., -2, 1:       ] = sval 
     314            ptab [..., -1, 1:       ] = sval    
     315            ptab [..., -1, 0        ] = sval 
    169316             
    170317        if cd_type == 'F' : 
    171             ptab[..., -2, 0:-1     ] = zero 
    172             ptab[..., -1, 0:-1     ] = zero 
    173             ptab[..., -1,  0       ] = zero 
    174             ptab[..., -1, -1       ] = zero 
     318            ptab [..., -2, 0:-1     ] = sval 
     319            ptab [..., -1, 0:-1     ] = sval 
     320            ptab [..., -1,  0       ] = sval 
     321            ptab [..., -1, -1       ] = sval 
    175322       
    176     if nperio in [5, 6] :            #  North fold F-point pivot   
     323    if nperio in [5, 6] :            #  North fold F-point pivot 
    177324        if cd_type in ['T', 'W']  : 
    178             ptab[..., -1, 0:       ] = zero 
     325            ptab [..., -1, 0:       ] = sval 
    179326                 
    180327        if cd_type == 'U' : 
    181             ptab[..., -1, 0:-1     ] = zero        
    182             ptab[..., -1, -1       ] = zero 
     328            ptab [..., -1, 0:-1     ] = sval        
     329            ptab [..., -1, -1       ] = sval 
    183330              
    184331        if cd_type == 'V' : 
    185             ptab[..., -1, 0:       ] = zero 
    186             ptab[..., -2, jpi//2:  ] = zero 
     332            ptab [..., -1, 0:       ] = sval 
     333            ptab [..., -2, jpi//2:  ] = sval 
    187334                              
    188335        if cd_type == 'F' : 
    189             ptab[..., -1, 0:-1     ] = zero 
    190             ptab[..., -1, -1       ] = zero 
    191             ptab[..., -2, jpi//2:-1] = zero 
     336            ptab [..., -1, 0:-1     ] = sval 
     337            ptab [..., -1, -1       ] = sval 
     338            ptab [..., -2, jpi//2:-1] = sval 
    192339 
    193340    return ptab 
     341 
    194342## =========================================================================== 
    195343## 
Note: See TracChangeset for help on using the changeset viewer.