source: TOOLS/MOSAIX/nemo.py @ 6064

Last change on this file since 6064 was 6064, checked in by omamce, 2 years ago

O.M. : changes in MOSAIX

  • Correct ocean mask to remove periodic duplicated points, to have full conservation of run-off
  • Suppress creation of corc mask, which is not used
  • Put SVN keywords in make_mosaic
  • Update nemo.py with somùe additionnal utilities
  • Adapt RunOffWeights?.py to new nemo module
  • Adapt CalvingWeights?.py to new ORCA configurations
  • Property svn:keywords set to Date Revision HeadURL Author Id
File size: 19.2 KB
Line 
1# -*- coding: utf-8 -*-
2## ===========================================================================
3##
4##  Warning, to install, configure, run, use any of Olivier Marti's
5##  software or to read the associated documentation you'll need at least
6##  one (1) brain in a reasonably working order. Lack of this implement
7##  will void any warranties (either express or implied).
8##  O. Marti assumes no responsability for errors, omissions,
9##  data loss, or any other consequences caused directly or indirectly by
10##  the usage of his software by incorrectly or partially configured
11##  personal.
12##
13'''
14Utilities to plot NEMO ORCA fields
15Periodicity and other stuff
16
17olivier.marti@lsce.ipsl.fr
18'''
19
20import sys, numpy as np
21try    : import xarray as xr
22except : pass
23
24rpi = np.pi ; rad = rpi / 180.0
25
26nperio_valid_range = [0,1,4,5,6]
27
28## SVN information
29__Author__   = "$Author$"
30__Date__     = "$Date$"
31__Revision__ = "$Revision$"
32__Id__       = "$Id$"
33__HeadURL    = "$HeadURL$"
34
35def __guessNperio__ (jpj, 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    if nperio == None :
44        if jpi ==  182 : nperio = 4 #   ORCA2. We choose legacy orca2.
45        if jpi ==  362 : nperio = 6 #   eORCA1.
46        if jpi == 1442 : nperio = 6 #   ORCA025.
47
48        if jpj ==  149 : nperio = 4 #   ORCA2. We choose legacy orca2.
49        if jpj ==  332 : nperio = 6 #   eORCA1.
50        #
51        if nperio == None :
52            sys.exit ('in nemo module : nperio not found, and cannot by guessed' )
53        else :
54            print ('nperio set as {:d} (deduced from jpi={:d})'.format (nperio, jpi))
55
56    return nperio
57
58def __guessPoint__ (ptab) :
59    '''
60    Tries to guess the grid point (periodicity parameter. See NEMO documentation for details)
61   
62    For array conforments with xgcm requirements
63
64    Inputs
65    ptab : xarray array
66
67    Credits : who is the original author ?
68    '''
69    gP = None
70    if isinstance (ptab, xr.core.dataarray.DataArray) :
71        if 'x_c' in ptab.dims and 'y_c' in ptab.dims                        : gP = 'T'
72        if 'x_f' in ptab.dims and 'y_c' in ptab.dims                        : gP = 'U'
73        if 'x_c' in ptab.dims and 'y_f' in ptab.dims                        : gP = 'V'
74        if 'x_f' in ptab.dims and 'y_f' in ptab.dims                        : gP = 'F'
75        if 'x_c' in ptab.dims and 'y_c' in ptab.dims and 'z_c' in ptab.dims : gP = 'T'
76        if 'x_c' in ptab.dims and 'y_c' in ptab.dims and 'z_f' in ptab.dims : gP = 'W'
77        if 'x_f' in ptab.dims and 'y_c' in ptab.dims and 'z_f' in ptab.dims : gP = 'U'
78        if 'x_c' in ptab.dims and 'y_f' in ptab.dims and 'z_f' in ptab.dims : gP = 'V'
79        if 'x_f' in ptab.dims and 'y_f' in ptab.dims and 'z_f' in ptab.dims : gP = 'F'
80             
81        if gP == None :
82            sys.exit ('in nemo module : cd_type not found, and cannot by guessed' )
83        else :
84            print ('Grid set as', gP, 'deduced from dims ', ptab.dims)
85            return gP
86    else :
87         sys.exit ('in nemo module : cd_type not found, input is not an xarray data' )
88         #return gP
89
90def fixed_lon (lon) :
91    '''
92    Returns corrected longitudes for nicer plots
93
94    fixed_lon (lon)
95    lon : longitudes of the grid. At least 2D.
96
97    Designed by Phil Pelson. See https://gist.github.com/pelson/79cf31ef324774c97ae7
98    '''
99    fixed_lon = lon.copy ()
100    for i, start in enumerate (np.argmax (np.abs (np.diff (lon, axis=-1)) > 180, axis=-1)) :
101        fixed_lon[..., i, start+1:] += 360
102
103    # Special case for eORCA025
104    if lon.shape[-1] == 1442 : lon [..., -2, :] =  lon [..., -3, :]
105       
106    return fixed_lon
107
108def jeq (lat) :
109    '''
110    Returns j index of equator in the grid
111   
112    lat : latitudes of the grid. At least 2D.
113    '''
114    jeq = np.argmin (np.abs (np.float64(lat[:,0])))
115    return jeq
116
117def lon1D (lon, lat=None) :
118    '''
119    Returns 1D longitude for simple plots.
120   
121    lon : longitudes of the grid
122    lat (optionnal) : latitudes of the grid
123    '''
124    if np.max (lat) != None :
125        je     = jeq (lat)
126        lon1D = lon[..., je, :]
127    else :
128        jpj, jpi = lon.shape[-2:]
129        lon1D = lon[..., jpj//3, :]
130
131    start = np.argmax (np.abs (np.diff (lon1D, axis=-1)) > 180.0, axis=-1)
132    lon1D[..., start+1:] += 360
133       
134    return lon1D
135
136def latreg (lat, diff=0.1) :
137    '''
138    Returns maximum j index where gridlines are along latitude in the northern hemisphere
139   
140    lat : latitudes of the grid (2D)
141    diff [optional] : tolerance
142    '''
143    if diff == None :
144        dy = np.float64 (np.mean (np.abs (lat - np.roll(lat,shift=1,axis=-2))))
145        diff = dy/100.
146   
147    je     = jeq (lat)
148    jreg   = np.where (lat[...,je:,:].max(axis=-1) - lat[...,je:,:].min(axis=-1)< diff)[-1][-1] + je
149    latreg = np.float64 (lat[...,jreg,:].mean(axis=-1))
150    JREG   = jreg
151
152    return jreg, latreg
153
154def lat1D (lat) :
155    '''
156    Returns 1D latitudes for zonal means and simple plots.
157
158    lat : latitudes of the grid (2D)
159    '''
160    jpj, jpi = lat.shape[-2:]
161    try :
162        if type (lat) == xr.core.dataarray.DataArray : math = xr
163        else : math = np
164    except : math = np
165
166    dy     = np.float64 (np.mean (np.abs (lat - np.roll (lat, shift=1,axis=-2))))
167    je     = jeq (lat)
168    lat_eq = np.float64 (lat[...,je,:].mean(axis=-1))
169     
170    jreg, lat_reg = latreg (lat)
171    lat_ave = np.mean (lat, axis=-1)
172
173    if ( np.abs (lat_eq) < dy/100. ) : # T, U or W grid
174        dys    = (90.-lat_reg) / (jpj-jreg-1)*0.5
175        yrange = 90.-dys-lat_reg
176    else                             :  # V or F grid
177        yrange = 90.    -lat_reg
178       
179    lat1D = math.where (lat_ave<lat_reg, lat_ave, lat_reg + yrange * (np.arange(jpj)-jreg)/(jpj-jreg-1) )   
180       
181    if math == xr : lat1D.attrs = lat.attrs
182
183    return lat1D
184
185def latlon1D (lat, lon) :
186    '''
187    Returns simple latitude and longitude (1D) for simple plots.
188
189    lat, lon : latitudes and longitudes of the grid (2D)
190    '''
191    return lat1D (lat),  lon1D (lon, lat)
192
193def mask_lonlat (ptab, x0, x1, y0, y1, lon, lat, sval=np.nan) :
194    try :
195        lon = lon.to_masked_array()
196        lat = lat.to_masked_array()
197    except : pass
198           
199    mask = np.logical_and ( np.logical_and(lat>y0, lat<y1), 
200        np.logical_or (np.logical_or (np.logical_and(lon>x0, lon<x1), np.logical_and(lon+360>x0, lon+360<x1)),
201                                         np.logical_and(lon-360>x0, lon-360<x1)) )
202    try    : tab = xr.where (mask, ptab, np.nan)
203    except : tab = np.where (mask, ptab, np.nan)
204   
205    return tab
206       
207def extend (tab, Lon=False, jplus=25, jpi=None, nperio=4) :
208    '''
209    Returns extended field eastward to have better plots, and box average crossing the boundary
210    Works only for xarray and numpy data (?)
211
212    tab : field to extend.
213    Lon : (optional, default=False) : if True, add 360 in the extended parts of the field
214    jpi : normal longitude dimension of the field. exrtend does nothing it the actual
215        size of the field != jpi (avoid to extend several times)
216    jplus (optional, default=25) : number of points added on the east side of the field
217   
218    '''
219    if tab.shape[-1] == 1 :
220        extend = tab
221
222    else :
223        if jpi == None : jpi = tab.shape[-1]
224   
225        try : 
226            if type (tab) ==  xr.core.dataarray.DataArray : math = xr
227            else : math = np
228        except : math = np
229
230        if Lon : xplus = -360.0
231        else   : xplus =    0.0
232
233        if tab.shape[-1] > jpi :
234            extend = tab
235        else :
236            if nperio == 1 :
237                istart = 0 ; le=jpi+1 ; la=0
238            if nperio > 1 : # OPA case with to halo points for periodicity
239                istart = 1 ; le=jpi-2 ; la=1  # Perfect, except at the pole that should be masked by lbc_plot
240               
241            if math == xr : 
242                extend = xr.concat      ( (tab[..., istart:istart+le+1] + xplus, tab[..., istart+la:jplus]), dim=tab.dims[-1] )
243            if math == np :
244                extend = np.concatenate ( (tab[..., istart:istart+le+1] + xplus, tab[..., istart+la:jplus]), axis=-1 )
245           
246    return extend
247
248def orca2reg (ff, lat_name='nav_lat', lon_name='nav_lon', y_name='y', x_name='x') :
249    '''
250    Assign ORCA dataset on a regular grid.
251    For use in the tropical region.
252
253    Inputs :
254      ff : xarray dataset
255      lat_name, lon_name : name of latitude and longitude 2D field in ff
256      y_name, x_name     : namex of dimensions in ff
257
258    Returns : xarray dataset with rectangular grid. Incorrect above 20°N
259    '''
260    # Compute 1D longitude and latitude
261    (lat, lon) = latlon1D (ff[lat_name], ff[lon_name])
262    # Assign lon and lat as dimensions of the dataset
263    if y_name in ff.dims : 
264        lat = xr.DataArray (lat, coords=[lat,], dims=['lat',])     
265        ff  = ff.rename_dims ({y_name: "lat",}).assign_coords (lat=lat)
266    if x_name in ff.dims :
267        lon = xr.DataArray (lon, coords=[lon,], dims=['lon',])
268        ff  = ff.rename_dims ({x_name: "lon",}).assign_coords (lon=lon)
269    # Force dimensions to be in the right order
270    coord_order = ['lat', 'lon']
271    for dim in [ 'depthw', 'depthv', 'depthu', 'deptht', 'depth', 'z',
272                 'time_counter', 'time', 'tbnds', 
273                 'bnds', 'axis_nbounds', 'two2', 'two1', 'two', 'four',] :
274        if dim in ff.dims : coord_order.insert (0, dim )
275       
276    ff = ff.transpose (*coord_order)
277    return ff
278
279def lbc_init (ptab, nperio=None, cd_type='T') :
280    """
281    Prepare for all lbc calls
282   
283    Set periodicity on input field
284    nperio    : Type of periodicity
285      0       : No periodicity
286      1, 4, 6 : Cyclic on i dimension (generaly longitudes) with 2 points halo
287      2       : Obsolete (was symmetric condition at southern boundary ?)
288      3, 4    : North fold T-point pivot (legacy ORCA2)
289      5, 6    : North fold F-point pivot (ORCA1, ORCA025, ORCA2 with new grid for paleo)
290    cd_type   : Grid specification : T, U, V or F
291
292    See NEMO documentation for further details
293    """
294    jpj, jpi = ptab.shape[-2:]
295    if nperio == None : nperio = __guessNperio__ (jpj, jpi, nperio)
296   
297    if nperio not in nperio_valid_range :
298        print ( 'nperio=', nperio, ' is not in the valid range', nperio_valid_range )
299        sys.exit()
300
301    if cd_type == None and isinstance (ptab, xr.core.dataarray.DataArray) :
302        cd_type = __guessPoint__ (ptab)
303
304    if cd_type not in ['T', 'U', 'V', 'F', 'W', 'F' ]:
305        print ( 'cd_type=', cd_type, ' is not in the list T, U, V, F, W, F' )
306        sys.exit ()
307
308    return jpj, jpi, nperio, cd_type
309       
310       
311def lbc (ptab, nperio=None, cd_type='T', psgn=1.0, nemo_4U_bug=False) :
312    """
313    Set periodicity on input field
314    ptab      : Input array (works
315      rank 2 at least : ptab[...., lat, lon]
316    nperio    : Type of periodicity
317    cd_type   : Grid specification : T, U, V or F
318    psgn      : For change of sign for vector components (1 for scalars, -1 for vector components)
319   
320    See NEMO documentation for further details
321    """
322    jpj, jpi, nperio, cd_type = lbc_init (ptab, nperio, cd_type)
323    psgn   = ptab.dtype.type (psgn)
324
325    #
326    #> East-West boundary conditions
327    # ------------------------------
328    if nperio in [1, 4, 6] :
329        # ... cyclic
330        ptab [..., :,  0] = ptab [..., :, -2]
331        ptab [..., :, -1] = ptab [..., :,  1]
332    #
333    #> North-South boundary conditions
334    # --------------------------------
335    if nperio in [3, 4] :  # North fold T-point pivot
336        if cd_type in [ 'T', 'W' ] : # T-, W-point
337            ptab [..., -1, 1:       ] = psgn * ptab [..., -3, -1:0:-1]
338            ptab [..., -1, 0        ] = psgn * ptab [..., -3, 2            ]
339            ptab [..., -2, jpi//2:  ] = psgn * ptab [..., -2, jpi//2:0:-1  ]
340               
341        if cd_type == 'U' :
342            ptab [..., -1, 0:-1     ] = psgn * ptab [..., -3, -1:0:-1      ]       
343            ptab [..., -1,  0       ] = psgn * ptab [..., -3,  1           ]
344            ptab [..., -1, -1       ] = psgn * ptab [..., -3, -2           ]
345               
346            if nemo_4U_bug :
347                ptab [..., -2, jpi//2+1:-1] = psgn * ptab [..., -2, jpi//2-2:0:-1]
348                ptab [..., -2, jpi//2-1   ] = psgn * ptab [..., -2, jpi//2     ]
349            else :
350                ptab [..., -2, jpi//2-1:-1] = psgn * ptab [..., -2, jpi//2:0:-1]
351               
352        if cd_type == 'V' : 
353            ptab [..., -2, 1:       ] = psgn * ptab [..., -3, jpi-1:0:-1   ]
354            ptab [..., -1, 1:       ] = psgn * ptab [..., -4, -1:0:-1      ]   
355            ptab [..., -1, 0        ] = psgn * ptab [..., -4, 2            ]
356               
357        if cd_type == 'F' :
358            ptab [..., -2, 0:-1     ] = psgn * ptab [..., -3, -1:0:-1      ]
359            ptab [..., -1, 0:-1     ] = psgn * ptab [..., -4, -1:0:-1      ]
360            ptab [..., -1,  0       ] = psgn * ptab [..., -4,  1           ]
361            ptab [..., -1, -1       ] = psgn * ptab [..., -4, -2           ] 
362       
363    if nperio in [5, 6] :            #  North fold F-point pivot 
364        if cd_type in ['T', 'W']  :
365            ptab [..., -1, 0:       ] = psgn * ptab [..., -2, -1::-1       ]
366               
367        if cd_type == 'U' :
368            ptab [..., -1, 0:-1     ] = psgn * ptab [..., -2, -2::-1       ]       
369            ptab [..., -1, -1       ] = psgn * ptab [..., -2, 0            ] # Bug ?
370               
371        if cd_type == 'V' :
372            ptab [..., -1, 0:       ] = psgn * ptab [..., -3, -1::-1       ]
373            ptab [..., -2, jpi//2:  ] = psgn * ptab [..., -2, jpi//2-1::-1 ]
374               
375        if cd_type == 'F' :
376            ptab [..., -1, 0:-1     ] = psgn * ptab [..., -3, -2::-1       ]
377            ptab [..., -1, -1       ] = psgn * ptab [..., -3, 0            ]
378            ptab [..., -2, jpi//2:-1] = psgn * ptab [..., -2, jpi//2-2::-1 ]
379               
380    return ptab
381
382def lbc_mask (ptab, nperio=None, cd_type='T', sval=np.nan) :
383    #
384    """
385    Mask fields on duplicated points
386    ptab      : Input array
387      rank 2 at least : ptab[...., lat, lon]
388    nperio    : Type of periodicity
389    cd_type   : Grid specification : T, U, V or F
390   
391    See NEMO documentation for further details
392    """
393
394    jpj, jpi, nperio, cd_type = lbc_init (ptab, nperio, cd_type)
395         
396    #
397    #> East-West boundary conditions
398    # ------------------------------
399    if nperio in [1, 4, 6] :
400        # ... cyclic
401        ptab [...,:,  0] = sval
402        ptab [...,:, -1] = sval
403
404    #
405    #> South (in which nperio cases ?)
406    # --------------------------------
407    if nperio in [1, 3, 4, 5, 6] :
408        ptab [...,0,:] = sval
409       
410    #
411    #> North-South boundary conditions
412    # --------------------------------
413    if nperio in [3, 4] :  # North fold T-point pivot
414        if cd_type in [ 'T', 'W' ] : # T-, W-point
415            ptab [..., -1, 1:       ] = sval
416            ptab [..., -1, 0        ] = sval
417            ptab [..., -2, jpi//2:  ] = sval
418               
419        if cd_type == 'U' :
420            ptab [..., -1,  :       ] = sval 
421            ptab [..., -2, jpi//2:  ] = sval
422               
423        if cd_type == 'V' :
424            ptab [..., -2, :       ] = sval
425            ptab [..., -1, :       ] = sval   
426               
427        if cd_type == 'F' :
428            ptab [..., -2, :     ] = sval
429            ptab [..., -1, :     ] = sval
430   
431    if nperio in [5, 6] :            #  North fold F-point pivot
432        if cd_type in ['T', 'W']  :
433            ptab [..., -1, 0:       ] = sval
434               
435        if cd_type == 'U' :
436            ptab [..., -1, 0:-1     ] = sval       
437            ptab [..., -1, -1       ] = sval
438             
439        if cd_type == 'V' :
440            ptab [..., -1, 0:       ] = sval
441            ptab [..., -2, jpi//2:  ] = sval
442                             
443        if cd_type == 'F' :
444            ptab [..., -1, 0:-1     ] = sval
445            ptab [..., -1, -1       ] = sval
446            ptab [..., -2, jpi//2:-1] = sval
447
448    return ptab
449
450
451def lbc_plot (ptab, nperio=None, cd_type='T', psgn=1.0, sval=np.nan) :
452    """
453    Set periodicity on input field, adapted for plotting for any cartopy projection
454    ptab      : Input array (works
455      rank 2 at least : ptab[...., lat, lon]
456    nperio    : Type of periodicity
457    cd_type   : Grid specification : T, U, V or F
458    psgn      : For change of sign for vector components (1 for scalars, -1 for vector components)
459   
460    See NEMO documentation for further details
461    """
462
463    jpj, jpi, nperio, cd_type = lbc_init (ptab, nperio, cd_type)
464    psgn   = ptab.dtype.type (psgn)
465   
466    #
467    #> East-West boundary conditions
468    # ------------------------------
469    if nperio in [1, 4, 6] :
470        # ... cyclic
471        ptab [..., :,  0] = ptab [..., :, -2]
472        ptab [..., :, -1] = ptab [..., :,  1]
473       
474    #
475    #> North-South boundary conditions
476    # --------------------------------
477    if nperio in [3, 4] :  # North fold T-point pivot
478        if cd_type in [ 'T', 'W' ] : # T-, W-point
479            ptab [..., -1, :] = sval
480            ptab [..., -2, :jpi//2] = sval
481            #ptab [..., -2, :] = sval
482                     
483        if cd_type == 'U' :
484            ptab [..., -1, : ] = sval
485
486        if cd_type == 'V' : 
487            ptab [..., -2, : ] = sval
488            ptab [..., -1, : ] = sval
489           
490        if cd_type == 'F' :
491            ptab [..., -2, : ] = sval
492            ptab [..., -1, : ] = sval
493     
494    if nperio in [5, 6] :            #  North fold F-point pivot 
495        if cd_type in ['T', 'W']  :
496            ptab [..., -1, : ] = sval
497               
498        if cd_type == 'U' :
499            ptab [..., -1, : ] = sval     
500             
501        if cd_type == 'V' :
502            ptab [..., -1, : ] = sval
503            ptab [..., -2, jpi//2:  ] = sval
504                             
505        if cd_type == 'F' :
506            ptab [..., -1, : ] = sval
507            ptab [..., -2, jpi//2:-1] = sval
508
509    return ptab
510
511def geo2oce (pxx, pyy, pzz, glam, gphi) : 
512    '''
513    Change vector from geocentric to east/north
514
515    Input
516        pxx, pyy, pzz : components on the geoce,tric system
517        glam, gphi : longitue and latitude of the points
518
519    '''
520    gsinlon = np.sin (rad * glam)
521    gcoslon = np.cos (rad * glam)
522    gsinlat = np.sin (rad * gphi)
523    gcoslat = np.cos (rad * gphi)
524         
525    pte = - pxx * gsinlon            + pyy * gcoslon
526    ptn = - pxx * gcoslon * gsinlat  - pyy * gsinlon * gsinlat + pzz * gcoslat
527
528    return pte, ptn
529
530def oce2geo (pte, ptn, glam, gphi) :
531    ##----------------------------------------------------------------------
532    ##                    ***  ROUTINE oce2geo  ***
533    ##     
534    ## ** Purpose :
535    ##
536    ## ** Method  :   Change vector from east/north to geocentric
537    ##
538    ## History :
539    ##        #         (A. Caubel)  oce2geo - Original code
540    ##----------------------------------------------------------------------
541    gsinlon = np.sin (rad * glam)
542    gcoslon = np.cos (rad * glam)
543    gsinlat = np.sin (rad * gphi)
544    gcoslat = np.cos (rad * gphi)
545
546    pxx = - pte * gsinlon - ptn * gcoslon * gsinlat
547    pyy =   pte * gcoslon - ptn * gsinlon * gsinlat
548    pzz =   ptn * gcoslat
549   
550    return pxx, pyy, pzz
551
552## ===========================================================================
553##
554##                               That's all folk's !!!
555##
556## ===========================================================================
Note: See TracBrowser for help on using the repository browser.