source: TOOLS/MOSAIX/Build_coordinates_mask.py @ 5947

Last change on this file since 5947 was 5947, checked in by snguyen, 3 years ago

Add argument parser to Build_coordinates_mask.py and lots of cleanup and cosmetics.
File size: 19.0 KB
Line 
1### ===========================================================================
2###
3### Build coordnates mask
4###
5### ===========================================================================
6# creates file coordinates_mask.nc used by mosaix from NEMO netcdf files
7# coordinates.nc and bathymetry.nc (on the same grid as coordinates.nc)
8
9import numpy  as np
10import xarray as xr
11import nemo
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
52
53# reproduce periodicity lbclnk bug for U and nperio = 4
54nemoUbug = myargs.nemo4Ubug
55# change metrics (and bathymetry) for straits
56straits = myargs.straits
57# periodicity imposed on coordinates, metrics and areas
58coordperio = myargs.coordPerio
59# function used to compute mask_T from bathymetry
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##
86
87# open input files while removing the time dimension
88f_coord = xr.open_dataset (n_coord, decode_times=False).squeeze()
89f_bathy = xr.open_dataset (n_bathy, decode_times=False).squeeze()
90
91# Suppress time if necessary
92try    :
93    del f_coord['time']
94    print ('time successfully removed')
95except :
96    pass
97    print ('failed to suppress time')
98
99# rename latitude, longitude and grid variables in bathymetry
100Bathymetry = f_bathy['Bathymetry'].copy()
101
102try :
103    Bathymetry = Bathymetry.rename ({'nav_lat':'nav_lat_grid_T', 'nav_lon':'nav_lon_grid_T'})
104    print ('Bathymetry file: Normal case')
105except : 
106    print ('Bathymetry file: Exception')
107    nav_lon_grid_T = f_bathy ['nav_lon']
108    nav_lat_grid_T = f_bathy ['nav_lat']
109    Bathymetry = xr.DataArray (Bathymetry, coords = { "nav_lat_grid_T": (["y", "x"], nav_lat_grid_T),
110                                                      "nav_lon_grid_T": (["y", "x"], nav_lon_grid_T) } )
111   
112Bathymetry = Bathymetry.rename ({'y':'y_grid_T', 'x':'x_grid_T'})
113
114# modify Bathymetry in straits for select grids (might change computed mask_T)
115if straits :
116    # Open straits for usual grids
117    if model == 'ORCA2.3' :
118        # orca_r2: Gibraltar strait open
119        Bathymetry[101,139] = 284.
120        # orca_r2: Bab el Mandeb strait open
121        Bathymetry[87,159] = 137.
122
123# impose periodicity of bathymetry
124Bathymetry = nemo.lbc (Bathymetry, nperio=nperio, cd_type='T')
125# suppress ocean points at southernmost position only if nperio in {3, 4, 5, 6}
126if nperio in (3, 4, 5, 6) :
127    Bathymetry[0,:] = 0.0
128
129##
130#!!! Create masks from bathymetry
131##
132
133# Creation of mask_T following choosed option maskbathy
134if maskbathynemo :
135    # Use same formula as domzgr.
136    mask_T = xr.where (Bathymetry - 1. + 0.1 >= 0.0, 1, 0).astype (dtype='f4')
137else :
138    mask_T = xr.where (Bathymetry > 0.0, 1, 0).astype (dtype='f4')
139
140# Creation of U, V, W, F masks from mask_T
141mask_U = mask_T * mask_T.shift (x_grid_T=-1)
142mask_V = mask_T * mask_T.shift (y_grid_T=-1)
143mask_F = mask_T * mask_T.shift (y_grid_T=-1) * mask_T.shift (x_grid_T=-1) * mask_T.shift (y_grid_T=-1, x_grid_T=-1)
144mask_W = mask_T
145
146# loop on TUVFW to modify coordinates and attributes of mask_[TUVFW] and maskutil_[TUVFW]
147for cd_type in ['T', 'U', 'V', 'F', 'W'] :
148    MaskName = 'mask_'     + cd_type
149    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]
159    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
163    locals()[UtilName].name = UtilName
164
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} )
190
191# remove _FillValue and missing_value
192        locals()[coord_name].encoding['_FillValue'] = None
193        locals()[coord_name].encoding['missing_value'] = None
194
195# define name attribute
196        locals()[coord_name].name = coord_name
197
198# add coordinates attributes
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##
211
212# create areas variables for grid cells from NEMO coordinates.nc file
213# create new variables e1 e2 to keep f_coord the same
214for cd_type in ['t', 'u', 'v', 'f'] :
215    for axis in ['1', '2'] :
216        coordName = 'e' + axis + cd_type
217        locals()[coordName]=f_coord[coordName].copy()
218        # remove zero values from areas
219        # need to be define for the extended grid south of -80S
220        # some point are undefined but you need to have e1 and e2 .NE. 0
221        locals()[coordName]=xr.where(locals()[coordName] == 0.0, 1.0e2, locals()[coordName])
222
223# Correct areas for straits
224if straits :
225    # ORCA R2 configuration
226    if model == 'ORCA2.3' :
227        # Gibraltar     : e2u reduced to 20 km
228        e2u[101,138:140] = 20.e3
229        # Bab el Mandeb : e2u reduced to 30 km
230        #                e1v reduced to 18 km
231        e1v[87,159] = 18.e3
232        e2u[87,159] = 30.e3
233        # Danish Straits: e2u reduced to 10 km
234        e2u[115,144:146] = 10.e3
235    # ORCA R1 configuration
236    if model == 'eORCA1.2' :
237        # Gibraltar : e2u reduced to 20 km
238        e2u[240,281:283] = 20.e3
239        # Bhosporus : e2u reduced to 10 km
240        e2u[247,313:315] =  10.e3
241        # Lombok : e1v reduced to 13 km
242        e1v[163:165,43] =  13.e3
243        # Sumba : e1v reduced to 8 km
244        e1v[163:165,47] =  8.e3
245        # Ombai : e1v reduced to 13 km
246        e1v[163:165,52] = 13.e3
247        # Timor Passage : e1v reduced to 20 km
248        #e1v[163:165,55] = 20.e3
249        # W Halmahera : e1v reduced to 30 km
250        e1v[180:182,54] = 30.e3
251        # E Halmahera : e1v reduced to 50 km
252        e1v[180:182,57] = 50.e3
253    # ORCA R05 configuration
254    if model == 'ORCA.05' :
255        # Reduced e2u at the Gibraltar Strait
256        e2u[326,562:564] =  20.e3
257        # Reduced e2u at the Bosphore Strait
258        e2u[342,626:628] =  10.e3
259        # Reduced e2u at the Sumba Strait
260        e2u[231,92:94] =  40.e3
261        # Reduced e2u at the Ombai Strait
262        e2u[231,102] =  15.e3
263        # Reduced e2u at the Palk Strait
264        e2u[269,14] =  10.e3
265        # Reduced e1v at the Lombok Strait
266        e1v[231:233,86] =  10.e3
267        # Reduced e1v at the Bab el Mandeb
268        e1v[275,661] =  25.e3
269
270# compute cells areas
271
272for cd_type in ['T', 'U', 'V', 'F', 'W'] :
273    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
289    locals()[areaName].encoding['_FillValue'] = None
290
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
388bounds_lon_grid_T, bounds_lat_grid_T = set_bounds ('T')
389bounds_lon_grid_U, bounds_lat_grid_U = set_bounds ('U')
390bounds_lon_grid_V, bounds_lat_grid_V = set_bounds ('V')
391bounds_lon_grid_W, bounds_lat_grid_W = set_bounds ('W')
392bounds_lon_grid_F, bounds_lat_grid_F = set_bounds ('F')
393
394# build xarray dataset to be saved
395ds = xr.Dataset ({
396    'mask_T'      : mask_T     ,
397    'mask_U'      : mask_U     ,
398    'mask_V'      : mask_V     ,
399    'mask_W'      : mask_W     ,
400    'mask_F'      : mask_F     ,
401    'area_grid_T' : area_grid_T,
402    'area_grid_U' : area_grid_U,
403    'area_grid_V' : area_grid_V,
404    'area_grid_W' : area_grid_W,
405    'area_grid_F' : area_grid_F,
406    'maskutil_T'  : maskutil_T ,
407    'maskutil_U'  : maskutil_U ,
408    'maskutil_V'  : maskutil_V ,
409    'maskutil_W'  : maskutil_W ,
410    'maskutil_F'  : maskutil_F ,
411    'bounds_lon_grid_T': bounds_lon_grid_T,
412    'bounds_lat_grid_T': bounds_lat_grid_T,
413    'bounds_lon_grid_U': bounds_lon_grid_U,
414    'bounds_lat_grid_U': bounds_lat_grid_U,
415    'bounds_lon_grid_V': bounds_lon_grid_V,
416    'bounds_lat_grid_V': bounds_lat_grid_V,
417    'bounds_lon_grid_W': bounds_lon_grid_W,
418    'bounds_lat_grid_W': bounds_lat_grid_W,
419    'bounds_lon_grid_F': bounds_lon_grid_F,
420    'bounds_lat_grid_F': bounds_lat_grid_F,
421})
422
423#replace nav_lon nav_lat with variables obtained from NEMO coordinates.nc file
424#by construction nav_lon nav_lat come from the bathymetry
425for cd_type in ['T', 'U', 'V', 'F', 'W'] :
426    for dir_type in ['lon', 'lat'] :
427        coord_name = 'nav_' + dir_type + '_grid_' + cd_type
428        ds.coords[coord_name]=locals()[coord_name]
429
430ds.attrs['name']         = 'coordinates_mask'
431ds.attrs['description']  = 'coordinates and mask for MOSAIX'
432ds.attrs['title']        = 'coordinates_mask'
433ds.attrs['source']       = 'IPSL Earth system model'
434ds.attrs['group']        = 'ICMC IPSL Climate Modelling Center'
435ds.attrs['Institution']  = 'IPSL https.//www.ipsl.fr'
436ds.attrs['Model']        = model
437ds.attrs['timeStamp']    = '{:%Y-%b-%d %H:%M:%S}'.format (datetime.datetime.now ())
438ds.attrs['history']      = 'Build from ' + n_coord + ' and ' + n_bathy
439ds.attrs['directory']    = os.getcwd     ()
440ds.attrs['user']         = os.getlogin   ()
441ds.attrs['HOSTNAME']     = platform.node ()
442ds.attrs['Python']       = 'Python version: ' +  platform.python_version ()
443ds.attrs['xarray']       = 'xarray version: ' +  xr.__version__
444ds.attrs['OS']           = platform.system  ()
445ds.attrs['release']      = platform.release ()
446ds.attrs['hardware']     = platform.machine ()
447
448ds.attrs['SVN_Author']   = "$Author:  $"
449ds.attrs['SVN_Date']     = "$Date:  $"
450ds.attrs['SVN_Revision'] = "$Revision:  $"
451ds.attrs['SVN_Id']       = "$Id: Build_coordinates_mask.py $"
452ds.attrs['SVN_HeadURL']  = "$HeadURL: http://forge.ipsl.jussieu.fr/igcmg/svn/TOOLS/MOSAIX/Build_coordinates_mask.py $"
453
454# save to output file
455ds.to_netcdf (n_out)
Note: See TracBrowser for help on using the repository browser.