Changeset 5947


Ignore:
Timestamp:
10/13/21 14:20:40 (2 years ago)
Author:
snguyen
Message:

Add argument parser to Build_coordinates_mask.py and lots of cleanup and cosmetics.

File:
1 edited

Legend:

Unmodified
Added
Removed
  • TOOLS/MOSAIX/Build_coordinates_mask.py

    r5942 r5947  
    1010import xarray as xr 
    1111import nemo 
    12 import datetime, os, platform 
     12import datetime, os, platform, argparse 
     13 
     14## SVN information 
     15__Author__   = "$Author: $" 
     16__Date__     = "$Date: $" 
     17__Revision__ = "$Revision: $" 
     18__Id__       = "$Id: $" 
     19__HeadURL__  = "$HeadURL: $" 
     20__SVN_Date__ = "$SVN_Date: $" 
     21### 
     22 
     23### ===== Handling command line parameters ================================================== 
     24# Creating a parser 
     25parser = argparse.ArgumentParser ( 
     26    description = """Read coordinates and bathymetry to build masks, grid bounds and areas for MOSAIX""", 
     27    epilog='-------- End of the help message --------') 
     28 
     29# Adding arguments 
     30parser.add_argument ('--model'   , help='oce model name', type=str, default='eORCA1.2', choices=('paleORCA', 'ORCA2.3', 'ORCA2.4', 'eORCA1.2', 'eORCA025', 'eORCA025.1') ) 
     31parser.add_argument ('--bathy'   , help='bathymetry file name', type=str, default='bathy_meter.nc' ) 
     32parser.add_argument ('--coord'   , help='coordinates file name', type=str, default='coordinates.nc' ) 
     33parser.add_argument ('--fout'    , help='output file name (given as an input to MOSAIX)', type=str, default='coordinates_mask.nc' ) 
     34parser.add_argument ('--ocePerio', help='periodicity of ocean grid', type=int, default=0, choices=(1, 2, 3, 4, 5, 6) ) 
     35parser.add_argument ('--nemo4Ubug',help='reproduce NEMO lbc_lnk bug for U grid and periodicity 4', type=bool, default=False, choices=(True, False) ) 
     36parser.add_argument ('--straits' , help='modify grid metrics for selected strait points (depends on model name)', type=bool, default=False, choices=(True, False) ) 
     37parser.add_argument ('--coordPerio', help='impose periodicity of coordinates and areas through lbc', type=bool, default=False, choices=(True, False) ) 
     38parser.add_argument ('--maskbathy', help='use the same formula as in NEMO to compute mask_T from the bathymetry file. Does not always give the same result as NEMO ???', type=bool, default=False, choices=(True, False) ) 
     39 
     40# Parse command line 
     41myargs = parser.parse_args() 
     42 
     43# ocean model name 
     44model = myargs.model 
     45 
     46# bathymetry input file 
     47n_bathy = myargs.bathy 
     48# coordinates input file 
     49n_coord = myargs.coord 
     50# coordinates and mask output file 
     51n_out = myargs.fout 
    1352 
    1453# reproduce periodicity lbclnk bug for U and nperio = 4 
    15 nemoUbug = False 
     54nemoUbug = myargs.nemo4Ubug 
    1655# change metrics (and bathymetry) for straits 
    17 straits = False 
     56straits = myargs.straits 
    1857# periodicity imposed on coordinates, metrics and areas 
    19 coordperio = False 
     58coordperio = myargs.coordPerio 
    2059# function used to compute mask_T from bathymetry 
    21 maskbathynemo = False 
    22  
    23 # type of grid treated : ORCA 1, 2, paleorca # to be improved 
    24  
    25 #model   = 'eORCA1.2' 
    26 #model   = 'ORCA2.3' 
    27 #model   = 'ORCA2.4' 
    28 model   = 'paleORCA' 
    29 #model   = 'ORCA.05' 
    30  
    31 if model == 'eORCA1.2' : 
    32     n_coord = 'eORCA1.2coordinates.nc' #'eORCA_R1_coordinates_v1.0.nc'  
    33     n_bathy = 'eORCA1.2bathy_meter.nc' #'eORCA_R1_bathy_meter_v2.2.nc' 
    34     #n_coord = 'https://vesg.ipsl.upmc.fr/thredds/dodsC/work/p86mart/GRAF/DATA/eORCA_R1_coordinates_v1.0.nc' 
    35     #n_bathy = 'https://vesg.ipsl.upmc.fr/thredds/dodsC/work/p86mart/GRAF/DATA/eORCA_R1_bathy_meter_v2.2.nc' 
    36     nperio  = 6 
    37     n_out   = 'eORCA1.2_coordinates_mask_test.nc' 
    38  
    39 if model == 'ORCA2.3' : 
    40     nperio  = 4 
    41     n_coord = 'coordinates_ORCA2.nc' # 'ORCA2.3_coordinates.nc' # 'coordinates_ORCA2_LIM3_PISCES.nc' # 
    42     n_bathy = 'bathy_meter_ORCA2.nc' # 'ORCA2.3_bathy_meter.nc' # 'bathy_meter_ORCA2_LIM3_PISCES.nc' # 
    43     n_out   = 'My_coordinates_xios_ORCA2.nc' 
    44      
    45 if model == 'ORCA2.4' : 
    46     nperio  = 4 
    47     n_coord = 'ORCA2.4_coordinates.nc' 
    48     n_bathy = 'ORCA2.4_bathy_meter.nc' 
    49     n_out   = 'ORCA2.4_coordinates_mask_test.nc' 
    50  
    51 if model == 'paleORCA' : 
    52     n_coord = 'coordinates_paleo.nc' 
    53     n_bathy = 'bathy_meter_paleo.nc' 
    54     nperio = 6 
    55     n_out = 'paleo_coordinates_mask_test.nc' 
    56  
    57 #--------------------------------------------------------------- 
    58 # function to generate grid bounds from NEMO coordinates.nc file 
    59 def set_bounds (cdgrd) : 
    60     ''' 
    61     Constructs lon/lat bounds 
    62     Bounds are numerated counter clockwise, from bottom left 
    63     See NEMO file OPA_SRC/IOM/iom.F90, ROUTINE set_grid_bounds, for more details 
    64     ''' 
    65     # Define offset of coordinate representing bottom-left corner 
    66     if cdgrd in ['T', 'W'] :  
    67         icnr = -1 ; jcnr = -1 
    68         corner_lon = f_coord ['glamf'].copy() ; corner_lat = f_coord ['gphif'].copy() 
    69         center_lon = f_coord ['glamt'].copy() ; center_lat = f_coord ['gphit'].copy() 
    70     if cdgrd == 'U'        :  
    71         icnr =  0 ; jcnr = -1 
    72         corner_lon = f_coord ['glamv'].copy() ; corner_lat = f_coord ['gphiv'].copy() 
    73         center_lon = f_coord ['glamu'].copy() ; center_lat = f_coord ['gphiu'].copy() 
    74     if cdgrd == 'V'        :  
    75         icnr = -1 ; jcnr =  0 
    76         corner_lon = f_coord ['glamu'].copy() ; corner_lat = f_coord ['gphiu'].copy() 
    77         center_lon = f_coord ['glamv'].copy() ; center_lat = f_coord ['gphiv'].copy() 
    78     if cdgrd == 'F'        :  
    79         icnr = -1 ; jcnr = -1 
    80         corner_lon = f_coord ['glamt'].copy() ; corner_lat = f_coord ['gphit'].copy() 
    81         center_lon = f_coord ['glamf'].copy() ; center_lat = f_coord ['gphif'].copy() 
    82        
    83     #y_grid, x_grid = corner_lon.shape ; 
    84     nvertex = 4 
    85     dims = ['y_grid_' + cdgrd, 'x_grid_' + cdgrd, 'nvertex_grid_' + cdgrd] 
    86      
    87     bounds_lon = xr.DataArray (np.zeros ((jpj, jpi, nvertex)), dims=dims) 
    88     bounds_lat = xr.DataArray (np.zeros ((jpj, jpi, nvertex)), dims=dims) 
    89      
    90     idx = [(jcnr,icnr), (jcnr,icnr+1), (jcnr+1,icnr+1), (jcnr+1,icnr)] 
    91      
    92     # Compute cell vertices that can be defined,  
    93     # and complete with periodicity 
    94     for nn in range (nvertex) : 
    95         tmp = np.roll (corner_lon, shift=tuple(-1*np.array(idx[nn])), axis=(-2,-1)) 
    96         bounds_lon[1:jpj,1:jpi,nn] = tmp[1:jpj,1:jpi] 
    97         tmp = np.roll (corner_lat, shift=tuple(-1*np.array(idx[nn])), axis=(-2,-1)) 
    98         bounds_lat[1:jpj,1:jpi,nn] = tmp[1:jpj,1:jpi] 
    99         bounds_lon[:,:,nn] = nemo.lbc (bounds_lon[:,:,nn], nperio=nperio, cd_type=cdgrd, nemo_4U_bug=nemoUbug) 
    100         bounds_lat[:,:,nn] = nemo.lbc (bounds_lat[:,:,nn], nperio=nperio, cd_type=cdgrd, nemo_4U_bug=nemoUbug) 
    101      
    102     # Zero-size cells at closed boundaries if cell points provided, 
    103     # otherwise they are closed cells with unrealistic bounds 
    104     if not (nperio == 1 or nperio == 4 or nperio == 6) : 
    105         for nn in range (nvertex) :  
    106             bounds_lon[:,0,nn]   = center_lon[:,0]   # (West or jpni = 1), closed E-W 
    107             bounds_lat[:,0,nn]   = center_lat[:,0] 
    108             bounds_lon[:,jpi,nn] = center_lon[:,jpi] # (East or jpni = 1), closed E-W 
    109             bounds_lat[:,jpi,nn] = center_lat[:,jpi] 
    110     if nperio != 2 : 
    111         for nn in range (nvertex) : 
    112             bounds_lon[0,:,nn] = center_lon[0,:] # (South or jpnj = 1), not symmetric 
    113             bounds_lat[0,:,nn] = center_lat[0,:] 
    114     if nperio < 3 : 
    115         for nn in range (nvertex) : 
    116             bounds_lon[jpj,:,nn] = center_lon[jpj,:] # (North or jpnj = 1), no north fold 
    117             bounds_lat[jpj,:,nn] = center_lat[jpj,:] 
    118      
    119     # Rotate cells at the north fold 
    120     if nperio >= 3 : 
    121         # Working array for location of northfold 
    122         z_fld = nemo.lbc (np.ones ((jpj, jpi)), nperio=nperio, cd_type=cdgrd, psgn=-1., nemo_4U_bug=nemoUbug) 
    123         z_fld = np.repeat((z_fld == -1.0)[...,np.newaxis],4,axis=2) 
    124         # circular shift of 2 indices in bounds third index 
    125         bounds_lon_tmp = np.roll (bounds_lon, shift=-2, axis=2) 
    126         bounds_lat_tmp = np.roll (bounds_lat, shift=-2, axis=2) 
    127         bounds_lon[:,:,:]  = np.where (z_fld, bounds_lon_tmp[:,:,:] , bounds_lon[:,:,:] ) 
    128         bounds_lat[:,:,:]  = np.where (z_fld, bounds_lat_tmp[:,:,:] , bounds_lat[:,:,:] ) 
    129      
    130     # Invert cells at the symmetric equator 
    131     if nperio == 2 : 
    132         bounds_lon_tmp = np.roll (bounds_lon, shift=-2, axis=2) 
    133         bounds_lat_tmp = np.roll (bounds_lat, shift=-2, axis=2) 
    134         bounds_lon [0,:,:] = bounds_lon [0,:,:] 
    135         bounds_lat [0,:,:] = bounds_lat [0,:,:] 
    136      
    137     #bounds_lon.attrs['coordinates'] = 'nav_lat_grid_' + cdgrd + ' nav_lon_grid_' + cdgrd 
    138     #bounds_lat.attrs['coordinates'] = 'nav_lat_grid_' + cdgrd + ' nav_lon_grid_' + cdgrd 
    139     bounds_lon.attrs['units']       = 'degrees_east' 
    140     bounds_lat.attrs['units']       = 'degrees_north' 
    141     bounds_lon.name = 'bounds_lon_grid_' + cdgrd 
    142     bounds_lat.name = 'bounds_lat_grid_' + cdgrd 
    143     # remove _FillValue 
    144     bounds_lon.encoding['_FillValue'] = None 
    145     bounds_lat.encoding['_FillValue'] = None 
    146  
    147     return bounds_lon, bounds_lat 
    148 #------------------------------------------------------------ 
     60maskbathynemo = myargs.maskbathy 
     61 
     62# type of grid periodicity 
     63nperio = myargs.ocePerio 
     64 
     65# check of periodicity with type of grid  
     66if model in ('eORCA1.2', 'paleORCA', 'ORCA025', 'eORCA025', 'eORCA025.1') : 
     67    if nperio != 0 : 
     68        if nperio != 6 : 
     69            print(f'Warning ! model = {model} and ocePerio = {nperio} instead of 6 !') 
     70    else : 
     71        nperio = 6 
     72 
     73if model in ('ORCA2.3', 'ORCA2.4') : 
     74    if nperio != 0 : 
     75        if nperio != 4 : 
     76            print(f'Attention ! model = {model} and ocePerio = {nperio} instead of 4 !') 
     77    else : 
     78        nperio = 4 
     79 
     80print(f' model = {model}\n ocePerio = {nperio}\n bathy = {n_bathy}\n coord = {n_coord}\n fout  = {n_out}') 
     81print(f' Switchs : nemo4Ubug = {nemoUbug}, straits = {straits}, coordPerio = {coordperio}, maskbathy = {maskbathynemo}') 
     82 
     83## 
     84#!!! bathymetry and coordinates files 
     85## 
    14986 
    15087# open input files while removing the time dimension 
     
    16097    print ('failed to suppress time') 
    16198 
     99# rename latitude, longitude and grid variables in bathymetry 
    162100Bathymetry = f_bathy['Bathymetry'].copy() 
    163101 
     
    174112Bathymetry = Bathymetry.rename ({'y':'y_grid_T', 'x':'x_grid_T'}) 
    175113 
     114# modify Bathymetry in straits for select grids (might change computed mask_T) 
    176115if straits : 
    177116    # Open straits for usual grids 
     
    182121        Bathymetry[87,159] = 137. 
    183122 
    184 # Determine periodicity from grid type 
    185 jpj, jpi = Bathymetry.shape 
    186 try :  
    187     if nperio != None : 
    188         print ('nperio specified : ', nperio) 
    189 except : 
    190     nperio = None 
    191     if jpi ==  182 : nperio = 4 # ORCA2. \!/ We choose legacy orca2 
    192     if jpi ==  362 : nperio = 6 # ORCA1. 
    193     if jpi == 1442 : nperio = 6 # ORCA025. 
    194     # 
    195     if nperio == None : 
    196         sys.exit ('nperio not found, and cannot by guessed' ) 
    197     else : 
    198         print ('nperio set as {:d} (deduced from data size {:d} {:d})'.format (nperio, jpj, jpi)) 
    199  
     123# impose periodicity of bathymetry 
    200124Bathymetry = nemo.lbc (Bathymetry, nperio=nperio, cd_type='T') 
    201 #Suppress ocean points at southernmost position only if nperio in {3, 4, 5, 6} 
     125# suppress ocean points at southernmost position only if nperio in {3, 4, 5, 6} 
    202126if nperio in (3, 4, 5, 6) : 
    203     Bathymetry[0,:] = 0.0  
    204  
    205 # Creation of mask_T 
    206  
     127    Bathymetry[0,:] = 0.0 
     128 
     129## 
     130#!!! Create masks from bathymetry 
     131## 
     132 
     133# Creation of mask_T following choosed option maskbathy 
    207134if maskbathynemo : 
    208135    # Use same formula as domzgr.  
    209     mask_T = xr.where (Bathymetry - 1. + 0.1 >= 0.0, 1, 0).astype (dtype='int8') 
     136    mask_T = xr.where (Bathymetry - 1. + 0.1 >= 0.0, 1, 0).astype (dtype='f4') 
    210137else : 
    211     mask_T = xr.where (Bathymetry > 0.0, 1, 0).astype (dtype='int8') 
     138    mask_T = xr.where (Bathymetry > 0.0, 1, 0).astype (dtype='f4') 
    212139 
    213140# Creation of U, V, W, F masks from mask_T 
     
    217144mask_W = mask_T 
    218145 
    219 mask_U = nemo.lbc (mask_U, nperio=nperio, cd_type='U', nemo_4U_bug=nemoUbug).astype (dtype='int8') 
    220 mask_V = nemo.lbc (mask_V, nperio=nperio, cd_type='V').astype (dtype='int8') 
    221 mask_F = nemo.lbc (mask_F, nperio=nperio, cd_type='F').astype (dtype='int8') 
    222 mask_W = nemo.lbc (mask_W, nperio=nperio, cd_type='W').astype (dtype='int8') 
    223  
    224 mask_U = mask_U.rename ( {'y_grid_T'      : 'y_grid_U'      , 'x_grid_T'      : 'x_grid_U',  
    225                           'nav_lat_grid_T': 'nav_lat_grid_U', 'nav_lon_grid_T': 'nav_lon_grid_U'} ) 
    226 mask_V = mask_V.rename ( {'y_grid_T'      : 'y_grid_V'      , 'x_grid_T'      : 'x_grid_V', 
    227                           'nav_lat_grid_T': 'nav_lat_grid_V', 'nav_lon_grid_T': 'nav_lon_grid_V'} ) 
    228 mask_W = mask_W.rename ( {'y_grid_T'      : 'y_grid_W'      , 'x_grid_T'      : 'x_grid_W', 
    229                           'nav_lat_grid_T': 'nav_lat_grid_W', 'nav_lon_grid_T': 'nav_lon_grid_W'} ) 
    230 mask_F = mask_F.rename ( {'y_grid_T'      : 'y_grid_F'      , 'x_grid_T'      : 'x_grid_F', 
    231                           'nav_lat_grid_T': 'nav_lat_grid_F', 'nav_lon_grid_T': 'nav_lon_grid_F'} ) 
    232  
    233 mask_T.name = 'mask_T' 
    234 mask_U.name = 'mask_U' 
    235 mask_V.name = 'mask_V' 
    236 mask_F.name = 'mask_F' 
    237 mask_W.name = 'mask_W' 
    238  
    239 # create masks without duplicate points : maskutil_[TUVWF] 
     146# loop on TUVFW to modify coordinates and attributes of mask_[TUVFW] and maskutil_[TUVFW] 
    240147for cd_type in ['T', 'U', 'V', 'F', 'W'] : 
    241148    MaskName = 'mask_'     + cd_type 
    242149    UtilName = 'maskutil_' + cd_type 
     150    # impose periodicity of chosen grid model on masks 
     151    locals()[MaskName] = nemo.lbc (locals()[MaskName], nperio=nperio, cd_type=cd_type, nemo_4U_bug=nemoUbug).astype (dtype='f4') 
     152    # rename masks coordinates 
     153    if cd_type != 'T' : 
     154        locals()[MaskName] = locals()[MaskName].rename \ 
     155                    ( {'y_grid_T'      :       'y_grid_'+cd_type, 'x_grid_T'      :       'x_grid_'+cd_type, 
     156                       'nav_lat_grid_T': 'nav_lat_grid_'+cd_type, 'nav_lon_grid_T': 'nav_lon_grid_'+cd_type} ) 
     157 
     158    # create masks without duplicate points : maskutil_[TUVWF] 
    243159    locals()[UtilName] = nemo.lbc_mask (locals()[MaskName].copy(), nperio=nperio, cd_type=cd_type) 
     160  
     161    #set name attribute of mask dataset 
     162    locals()[MaskName].name = MaskName 
    244163    locals()[UtilName].name = UtilName 
    245164 
    246 # add masks attributes 
    247 mask_T.attrs    ['cell_measures'] = 'area: area_grid_T' 
    248 mask_U.attrs    ['cell_measures'] = 'area: area_grid_U' 
    249 mask_V.attrs    ['cell_measures'] = 'area: area_grid_V' 
    250 mask_W.attrs    ['cell_measures'] = 'area: area_grid_W' 
    251 mask_F.attrs    ['cell_measures'] = 'area: area_grid_F' 
    252  
    253 maskutil_T.attrs['cell_measures'] = 'area: area_grid_T' 
    254 maskutil_U.attrs['cell_measures'] = 'area: area_grid_U' 
    255 maskutil_V.attrs['cell_measures'] = 'area: area_grid_V' 
    256 maskutil_W.attrs['cell_measures'] = 'area: area_grid_W' 
    257 maskutil_F.attrs['cell_measures'] = 'area: area_grid_F' 
    258  
    259 # create grid coordinates from NEMO coordinates.nc file 
    260 if coordperio : 
    261     nav_lon_grid_T = nemo.lbc (f_coord ['glamt'].copy(), nperio=nperio, cd_type='T') 
    262     nav_lat_grid_T = nemo.lbc (f_coord ['gphit'].copy(), nperio=nperio, cd_type='T') 
    263     nav_lon_grid_U = nemo.lbc (f_coord ['glamu'].copy(), nperio=nperio, cd_type='U',nemo_4U_bug=nemoUbug) 
    264     nav_lat_grid_U = nemo.lbc (f_coord ['gphiu'].copy(), nperio=nperio, cd_type='U',nemo_4U_bug=nemoUbug) 
    265     nav_lon_grid_V = nemo.lbc (f_coord ['glamv'].copy(), nperio=nperio, cd_type='V') 
    266     nav_lat_grid_V = nemo.lbc (f_coord ['gphiv'].copy(), nperio=nperio, cd_type='V') 
    267     nav_lon_grid_W = nemo.lbc (f_coord ['glamt'].copy(), nperio=nperio, cd_type='W') 
    268     nav_lat_grid_W = nemo.lbc (f_coord ['gphit'].copy(), nperio=nperio, cd_type='W') 
    269     nav_lon_grid_F = nemo.lbc (f_coord ['glamf'].copy(), nperio=nperio, cd_type='F') 
    270     nav_lat_grid_F = nemo.lbc (f_coord ['gphif'].copy(), nperio=nperio, cd_type='F') 
    271 else : 
    272     nav_lon_grid_T = f_coord ['glamt'].copy() 
    273     nav_lat_grid_T = f_coord ['gphit'].copy() 
    274     nav_lon_grid_U = f_coord ['glamu'].copy() 
    275     nav_lat_grid_U = f_coord ['gphiu'].copy() 
    276     nav_lon_grid_V = f_coord ['glamv'].copy() 
    277     nav_lat_grid_V = f_coord ['gphiv'].copy() 
    278     nav_lon_grid_W = f_coord ['glamt'].copy() 
    279     nav_lat_grid_W = f_coord ['gphit'].copy() 
    280     nav_lon_grid_F = f_coord ['glamf'].copy() 
    281     nav_lat_grid_F = f_coord ['gphif'].copy() 
    282  
    283 nav_lon_grid_T = nav_lon_grid_T.rename( {'y':'y_grid_T', 'x':'x_grid_T'} ) 
    284 nav_lat_grid_T = nav_lat_grid_T.rename( {'y':'y_grid_T', 'x':'x_grid_T'} ) 
    285 nav_lon_grid_U = nav_lon_grid_U.rename( {'y':'y_grid_U', 'x':'x_grid_U'} ) 
    286 nav_lat_grid_U = nav_lat_grid_U.rename( {'y':'y_grid_U', 'x':'x_grid_U'} ) 
    287 nav_lon_grid_V = nav_lon_grid_V.rename( {'y':'y_grid_V', 'x':'x_grid_V'} ) 
    288 nav_lat_grid_V = nav_lat_grid_V.rename( {'y':'y_grid_V', 'x':'x_grid_V'} ) 
    289 nav_lon_grid_W = nav_lon_grid_W.rename( {'y':'y_grid_W', 'x':'x_grid_W'} ) 
    290 nav_lat_grid_W = nav_lat_grid_W.rename( {'y':'y_grid_W', 'x':'x_grid_W'} ) 
    291 nav_lon_grid_F = nav_lon_grid_F.rename( {'y':'y_grid_F', 'x':'x_grid_F'} ) 
    292 nav_lat_grid_F = nav_lat_grid_F.rename( {'y':'y_grid_F', 'x':'x_grid_F'} ) 
     165    # remove _FillVallue from mskutil 
     166    locals()[MaskName].encoding['_FillValue'] = None 
     167    locals()[UtilName].encoding['_FillValue'] = None 
     168 
     169    # add masks attributes 
     170    locals()[MaskName].attrs['cell_measures'] = 'area: area_grid_'+cd_type 
     171    locals()[UtilName].attrs['cell_measures'] = 'area: area_grid_'+cd_type 
     172 
     173## 
     174#!!! create grid coordinates from NEMO coordinates.nc file 
     175## 
     176 
     177angle = { 'lon' : 'glam', 'lat' : 'gphi' } 
     178gridv = { 'T' : 't', 'U' : 'u', 'V' : 'v', 'F' : 'f', 'W' : 't' } 
     179for cd_type in ['T', 'U', 'V', 'F', 'W'] : 
     180    for dir_type in ['lon', 'lat'] : 
     181        coord_name = 'nav_' + dir_type + '_grid_' + cd_type 
     182        dir_name = angle[dir_type]+gridv[cd_type] 
     183# impose or not periodicity on coordinates read from file 
     184        if coordperio : 
     185            locals()[coord_name] = nemo.lbc (f_coord [dir_name].copy(), nperio=nperio, cd_type=cd_type, nemo_4U_bug=nemoUbug) 
     186        else : 
     187            locals()[coord_name] = f_coord [dir_name].copy() 
     188 
     189        locals()[coord_name] = locals()[coord_name].rename( {'y':'y_grid_'+cd_type, 'x':'x_grid_'+cd_type} ) 
    293190 
    294191# remove _FillValue and missing_value 
    295 for cd_type in ['T', 'U', 'V', 'F', 'W'] : 
    296     for dir_type in ['lat', 'lon'] : 
    297         coord_name = 'nav_' + dir_type + '_grid_' + cd_type 
    298192        locals()[coord_name].encoding['_FillValue'] = None 
    299193        locals()[coord_name].encoding['missing_value'] = None 
    300194 
    301 nav_lon_grid_T.name = 'nav_lon_grid_T' 
    302 nav_lat_grid_T.name = 'nav_lat_grid_T' 
    303 nav_lon_grid_U.name = 'nav_lon_grid_U' 
    304 nav_lat_grid_U.name = 'nav_lat_grid_U' 
    305 nav_lon_grid_V.name = 'nav_lon_grid_V' 
    306 nav_lat_grid_V.name = 'nav_lat_grid_V' 
    307 nav_lon_grid_F.name = 'nav_lon_grid_F' 
    308 nav_lat_grid_F.name = 'nav_lat_grid_F' 
    309 nav_lon_grid_W.name = 'nav_lon_grid_W' 
    310 nav_lat_grid_W.name = 'nav_lat_grid_W' 
     195# define name attribute 
     196        locals()[coord_name].name = coord_name 
    311197 
    312198# add coordinates attributes 
    313 nav_lon_grid_T.attrs['standard_name'] = 'longitude' 
    314 nav_lon_grid_T.attrs['long_name']     = 'Longitude' 
    315 nav_lon_grid_T.attrs['units']         = 'degrees_east' 
    316 nav_lat_grid_T.attrs['standard_name'] = 'latitude' 
    317 nav_lat_grid_T.attrs['long_name']     = 'Latitude' 
    318 nav_lat_grid_T.attrs['units']         = 'degrees_north' 
    319  
    320 nav_lon_grid_U.attrs['standard_name'] = 'longitude' 
    321 nav_lon_grid_U.attrs['long_name']     = 'Longitude' 
    322 nav_lon_grid_U.attrs['units']         = 'degrees_east' 
    323 nav_lat_grid_U.attrs['standard_name'] = 'latitude' 
    324 nav_lat_grid_U.attrs['long_name']     = 'Latitude' 
    325 nav_lat_grid_U.attrs['units']         = 'degrees_north' 
    326  
    327 nav_lon_grid_V.attrs['standard_name'] = 'longitude' 
    328 nav_lon_grid_V.attrs['long_name']     = 'Longitude' 
    329 nav_lon_grid_V.attrs['units']         = 'degrees_east' 
    330 nav_lat_grid_V.attrs['standard_name'] = 'latitude' 
    331 nav_lat_grid_V.attrs['long_name']     = 'Latitude' 
    332 nav_lat_grid_V.attrs['units']         = 'degrees_north' 
    333  
    334 nav_lon_grid_W.attrs['standard_name'] = 'longitude' 
    335 nav_lon_grid_W.attrs['long_name']     = 'Longitude' 
    336 nav_lon_grid_W.attrs['units']         = 'degrees_east' 
    337 nav_lat_grid_W.attrs['standard_name'] = 'latitude' 
    338 nav_lat_grid_W.attrs['long_name']     = 'Latitude' 
    339 nav_lat_grid_W.attrs['units']         = 'degrees_north' 
    340  
    341 nav_lon_grid_F.attrs['standard_name'] = 'longitude' 
    342 nav_lon_grid_F.attrs['long_name']     = 'Longitude' 
    343 nav_lon_grid_F.attrs['units']         = 'degrees_east' 
    344 nav_lat_grid_F.attrs['standard_name'] = 'latitude' 
    345 nav_lat_grid_F.attrs['long_name']     = 'Latitude' 
    346 nav_lat_grid_F.attrs['units']         = 'degrees_north' 
    347  
    348 nav_lon_grid_T.attrs['bounds'] = 'bounds_lon_grid_T' 
    349 nav_lat_grid_T.attrs['bounds'] = 'bounds_lat_grid_T' 
    350 nav_lon_grid_U.attrs['bounds'] = 'bounds_lon_grid_U' 
    351 nav_lat_grid_U.attrs['bounds'] = 'bounds_lat_grid_U' 
    352 nav_lon_grid_V.attrs['bounds'] = 'bounds_lon_grid_V' 
    353 nav_lat_grid_V.attrs['bounds'] = 'bounds_lat_grid_V' 
    354 nav_lon_grid_W.attrs['bounds'] = 'bounds_lon_grid_W' 
    355 nav_lat_grid_W.attrs['bounds'] = 'bounds_lat_grid_W' 
    356 nav_lon_grid_F.attrs['bounds'] = 'bounds_lon_grid_F' 
    357 nav_lat_grid_F.attrs['bounds'] = 'bounds_lat_grid_F' 
     199        locals()[coord_name].attrs['bounds'] ='bounds_' + dir_type + '_grid_' + cd_type 
     200 
     201    locals()['nav_lon_grid_'+cd_type].attrs['standard_name'] = 'longitude' 
     202    locals()['nav_lon_grid_'+cd_type].attrs['long_name']     = 'Longitude' 
     203    locals()['nav_lon_grid_'+cd_type].attrs['units']         = 'degrees_east' 
     204    locals()['nav_lat_grid_'+cd_type].attrs['standard_name'] = 'latitude' 
     205    locals()['nav_lat_grid_'+cd_type].attrs['long_name']     = 'Latitude' 
     206    locals()['nav_lat_grid_'+cd_type].attrs['units']         = 'degrees_north' 
     207 
     208## 
     209#!!! compute areas of cells at coordinate points 
     210## 
    358211 
    359212# create areas variables for grid cells from NEMO coordinates.nc file 
     
    367220        # some point are undefined but you need to have e1 and e2 .NE. 0 
    368221        locals()[coordName]=xr.where(locals()[coordName] == 0.0, 1.0e2, locals()[coordName]) 
    369  
    370 #Compute cells areas 
    371222 
    372223# Correct areas for straits 
     
    417268        e1v[275,661] =  25.e3 
    418269 
    419 if coordperio : 
    420     area_grid_T = nemo.lbc (e1t*e2t, nperio=nperio, cd_type='T') 
    421     area_grid_U = nemo.lbc (e1u*e2u, nperio=nperio, cd_type='U',nemo_4U_bug=nemoUbug) 
    422     area_grid_V = nemo.lbc (e1v*e2v, nperio=nperio, cd_type='V') 
    423     area_grid_W = nemo.lbc (e1t*e2t, nperio=nperio, cd_type='W') 
    424     area_grid_F = nemo.lbc (e1f*e2f, nperio=nperio, cd_type='F') 
    425 else : 
    426     area_grid_T = e1t*e2t 
    427     area_grid_U = e1u*e2u 
    428     area_grid_V = e1v*e2v 
    429     area_grid_W = e1t*e2t 
    430     area_grid_F = e1f*e2f 
    431  
    432 area_grid_T.name = 'area_grid_T' 
    433 area_grid_U.name = 'area_grid_U' 
    434 area_grid_V.name = 'area_grid_V' 
    435 area_grid_F.name = 'area_grid_F' 
    436 area_grid_W.name = 'area_grid_W' 
    437  
    438 area_grid_T = area_grid_T.rename ({'y':'y_grid_T', 'x':'x_grid_T'}) 
    439 area_grid_U = area_grid_U.rename ({'y':'y_grid_U', 'x':'x_grid_U'}) 
    440 area_grid_V = area_grid_V.rename ({'y':'y_grid_V', 'x':'x_grid_V'}) 
    441 area_grid_W = area_grid_W.rename ({'y':'y_grid_W', 'x':'x_grid_W'}) 
    442 area_grid_F = area_grid_F.rename ({'y':'y_grid_F', 'x':'x_grid_F'}) 
    443  
    444 area_grid_T.attrs['standard_name'] = 'cell_area' 
    445 area_grid_T.attrs['units']         = 'm2' 
    446 area_grid_U.attrs['standard_name'] = 'cell_area' 
    447 area_grid_U.attrs['units']         = 'm2' 
    448 area_grid_V.attrs['standard_name'] = 'cell_area' 
    449 area_grid_V.attrs['units']         = 'm2' 
    450 area_grid_W.attrs['standard_name'] = 'cell_area' 
    451 area_grid_W.attrs['units']         = 'm2' 
    452 area_grid_F.attrs['standard_name'] = 'cell_area' 
    453 area_grid_F.attrs['units']         = 'm2' 
     270# compute cells areas 
    454271 
    455272for cd_type in ['T', 'U', 'V', 'F', 'W'] : 
    456273    areaName = 'area_grid_' + cd_type 
     274    if coordperio : 
     275        locals()[areaName] = nemo.lbc (locals()['e1'+gridv[cd_type]]*locals()['e2'+gridv[cd_type]], 
     276                                                           nperio=nperio, cd_type=cd_type, nemo_4U_bug=nemoUbug) 
     277    else : 
     278        locals()[areaName] = locals()['e1'+gridv[cd_type]]*locals()['e2'+gridv[cd_type]] 
     279 
     280    # rename indices 
     281    locals()[areaName] = locals()[areaName].rename ({'y':'y_grid_'+cd_type, 'x':'x_grid_'+cd_type}) 
     282 
     283    # add attributes 
     284    locals()[areaName].name = areaName 
     285    locals()[areaName].attrs['standard_name'] = 'cell_area' 
     286    locals()[areaName].attrs['units']         = 'm2' 
     287 
     288    # remove fill values 
    457289    locals()[areaName].encoding['_FillValue'] = None 
    458290 
    459 # compute grid cells bounds 
     291## 
     292#!!! compute grid bounds !!! 
     293## 
     294 
     295#--------------------------------------------------------------- 
     296# function to generate grid bounds from NEMO coordinates.nc file 
     297def set_bounds (cdgrd) : 
     298    ''' 
     299    Constructs lon/lat bounds 
     300    Bounds are numerated counter clockwise, from bottom left 
     301    See NEMO file OPA_SRC/IOM/iom.F90, ROUTINE set_grid_bounds, for more details 
     302    ''' 
     303    # Define offset of coordinate representing bottom-left corner 
     304    if cdgrd in ['T', 'W'] :  
     305        icnr = -1 ; jcnr = -1 
     306        corner_lon = f_coord ['glamf'].copy() ; corner_lat = f_coord ['gphif'].copy() 
     307        center_lon = f_coord ['glamt'].copy() ; center_lat = f_coord ['gphit'].copy() 
     308    if cdgrd == 'U'        :  
     309        icnr =  0 ; jcnr = -1 
     310        corner_lon = f_coord ['glamv'].copy() ; corner_lat = f_coord ['gphiv'].copy() 
     311        center_lon = f_coord ['glamu'].copy() ; center_lat = f_coord ['gphiu'].copy() 
     312    if cdgrd == 'V'        :  
     313        icnr = -1 ; jcnr =  0 
     314        corner_lon = f_coord ['glamu'].copy() ; corner_lat = f_coord ['gphiu'].copy() 
     315        center_lon = f_coord ['glamv'].copy() ; center_lat = f_coord ['gphiv'].copy() 
     316    if cdgrd == 'F'        :  
     317        icnr = -1 ; jcnr = -1 
     318        corner_lon = f_coord ['glamt'].copy() ; corner_lat = f_coord ['gphit'].copy() 
     319        center_lon = f_coord ['glamf'].copy() ; center_lat = f_coord ['gphif'].copy() 
     320       
     321    jpj, jpi = corner_lon.shape ; 
     322    nvertex = 4 
     323    dims = ['y_grid_' + cdgrd, 'x_grid_' + cdgrd, 'nvertex_grid_' + cdgrd] 
     324     
     325    bounds_lon = xr.DataArray (np.zeros ((jpj, jpi, nvertex)), dims=dims) 
     326    bounds_lat = xr.DataArray (np.zeros ((jpj, jpi, nvertex)), dims=dims) 
     327     
     328    idx = [(jcnr,icnr), (jcnr,icnr+1), (jcnr+1,icnr+1), (jcnr+1,icnr)] 
     329     
     330    # Compute cell vertices that can be defined,  
     331    # and complete with periodicity 
     332    for nn in range (nvertex) : 
     333        tmp = np.roll (corner_lon, shift=tuple(-1*np.array(idx[nn])), axis=(-2,-1)) 
     334        bounds_lon[1:jpj,1:jpi,nn] = tmp[1:jpj,1:jpi] 
     335        tmp = np.roll (corner_lat, shift=tuple(-1*np.array(idx[nn])), axis=(-2,-1)) 
     336        bounds_lat[1:jpj,1:jpi,nn] = tmp[1:jpj,1:jpi] 
     337        bounds_lon[:,:,nn] = nemo.lbc (bounds_lon[:,:,nn], nperio=nperio, cd_type=cdgrd, nemo_4U_bug=nemoUbug) 
     338        bounds_lat[:,:,nn] = nemo.lbc (bounds_lat[:,:,nn], nperio=nperio, cd_type=cdgrd, nemo_4U_bug=nemoUbug) 
     339     
     340    # Zero-size cells at closed boundaries if cell points provided, 
     341    # otherwise they are closed cells with unrealistic bounds 
     342    if not (nperio == 1 or nperio == 4 or nperio == 6) : 
     343        for nn in range (nvertex) :  
     344            bounds_lon[:,0,nn]   = center_lon[:,0]   # (West or jpni = 1), closed E-W 
     345            bounds_lat[:,0,nn]   = center_lat[:,0] 
     346            bounds_lon[:,jpi,nn] = center_lon[:,jpi] # (East or jpni = 1), closed E-W 
     347            bounds_lat[:,jpi,nn] = center_lat[:,jpi] 
     348    if nperio != 2 : 
     349        for nn in range (nvertex) : 
     350            bounds_lon[0,:,nn] = center_lon[0,:] # (South or jpnj = 1), not symmetric 
     351            bounds_lat[0,:,nn] = center_lat[0,:] 
     352    if nperio < 3 : 
     353        for nn in range (nvertex) : 
     354            bounds_lon[jpj,:,nn] = center_lon[jpj,:] # (North or jpnj = 1), no north fold 
     355            bounds_lat[jpj,:,nn] = center_lat[jpj,:] 
     356     
     357    # Rotate cells at the north fold 
     358    if nperio >= 3 : 
     359        # Working array for location of northfold 
     360        z_fld = nemo.lbc (np.ones ((jpj, jpi)), nperio=nperio, cd_type=cdgrd, psgn=-1., nemo_4U_bug=nemoUbug) 
     361        z_fld = np.repeat((z_fld == -1.0)[...,np.newaxis],4,axis=2) 
     362        # circular shift of 2 indices in bounds third index 
     363        bounds_lon_tmp = np.roll (bounds_lon, shift=-2, axis=2) 
     364        bounds_lat_tmp = np.roll (bounds_lat, shift=-2, axis=2) 
     365        bounds_lon[:,:,:]  = np.where (z_fld, bounds_lon_tmp[:,:,:] , bounds_lon[:,:,:] ) 
     366        bounds_lat[:,:,:]  = np.where (z_fld, bounds_lat_tmp[:,:,:] , bounds_lat[:,:,:] ) 
     367     
     368    # Invert cells at the symmetric equator 
     369    if nperio == 2 : 
     370        bounds_lon_tmp = np.roll (bounds_lon, shift=-2, axis=2) 
     371        bounds_lat_tmp = np.roll (bounds_lat, shift=-2, axis=2) 
     372        bounds_lon [0,:,:] = bounds_lon [0,:,:] 
     373        bounds_lat [0,:,:] = bounds_lat [0,:,:] 
     374     
     375    #bounds_lon.attrs['coordinates'] = 'nav_lat_grid_' + cdgrd + ' nav_lon_grid_' + cdgrd 
     376    #bounds_lat.attrs['coordinates'] = 'nav_lat_grid_' + cdgrd + ' nav_lon_grid_' + cdgrd 
     377    bounds_lon.attrs['units']       = 'degrees_east' 
     378    bounds_lat.attrs['units']       = 'degrees_north' 
     379    bounds_lon.name = 'bounds_lon_grid_' + cdgrd 
     380    bounds_lat.name = 'bounds_lat_grid_' + cdgrd 
     381    # remove _FillValue 
     382    bounds_lon.encoding['_FillValue'] = None 
     383    bounds_lat.encoding['_FillValue'] = None 
     384 
     385    return bounds_lon, bounds_lat 
     386#------------------------------------------------------------ 
     387 
    460388bounds_lon_grid_T, bounds_lat_grid_T = set_bounds ('T') 
    461389bounds_lon_grid_U, bounds_lat_grid_U = set_bounds ('U') 
     
    494422 
    495423#replace nav_lon nav_lat with variables obtained from NEMO coordinates.nc file 
    496 #by construction nav_lon nav_lat are those of the bathymetry 
     424#by construction nav_lon nav_lat come from the bathymetry 
    497425for cd_type in ['T', 'U', 'V', 'F', 'W'] : 
    498     for dir_type in ['lat', 'lon'] : 
     426    for dir_type in ['lon', 'lat'] : 
    499427        coord_name = 'nav_' + dir_type + '_grid_' + cd_type 
    500428        ds.coords[coord_name]=locals()[coord_name] 
Note: See TracChangeset for help on using the changeset viewer.