source: TOOLS/MOSAIX/nemo.py @ 6041

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

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.
  • Property svn:keywords set to Date Revision HeadURL Author Id
File size: 11.9 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
21
22try    : import numpy as np
23except : pass
24
25try    : import xarray as xr
26except : pass
27
28## SVN information
29__Author__   = "$Author$"
30__Date__     = "$Date$"
31__Revision__ = "$Revision$"
32__Id__       = "$Id$"
33__HeadURL    = "$HeadURL$"
34
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   
44    if nperio == None :
45        if jpi ==  182 : nperio = 4 #   ORCA2. We choose legacy orca2.
46        if jpi ==  362 : nperio = 6 #   ORCA1.
47        if jpi == 1442 : nperio = 6 # ORCA025.
48        #
49        if nperio == None :
50            sys.exit ('in nemo.lbc : nperio not found, and cannot by guessed' )
51        else :
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) :
186    """
187    Set periodicity on input field
188    ptab      : Input array (works
189      rank 2 at least : ptab[...., lat, lon]
190    nperio    : Type of periodicity
191      1, 4, 6 : Cyclic on i dimension (generaly longitudes)
192      2       : Obsolete (was symmetric condition at southern boundary ?)
193      3, 4    : North fold T-point pivot (legacy ORCA2)
194      5, 6    : North fold F-point pivot (ORCA1, ORCA025, ORCA2 with new grid for paleo)
195    cd_type   : Grid specification : T, U, V or F
196    psgn      : For change of sign for vector components (1 for scalars, -1 for vector components)
197   
198    See NEMO documentation for further details
199    """
200
201    jpi    = ptab.shape[-1]
202    nperio = __guessNperio__ ( jpi, nperio)
203    psgn   = ptab.dtype.type(psgn)
204   
205    #
206    #> East-West boundary conditions
207    # ------------------------------
208    if nperio in [1, 4, 6] :
209        # ... cyclic
210        ptab [..., :,  0] = ptab [..., :, -2]
211        ptab [..., :, -1] = ptab [..., :,  1]
212   
213    #
214    #> North-South boundary conditions
215    # --------------------------------
216    if nperio in [3, 4] :  # North fold T-point pivot
217        if cd_type in [ 'T', 'W' ] : # T-, W-point
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  ]
221                     
222        if cd_type == 'U' :
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]
232           
233        if cd_type == 'V' : 
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            ]
237           
238        if cd_type == 'F' :
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           ] 
243     
244    if nperio in [5, 6] :            #  North fold F-point pivot 
245        if cd_type in ['T', 'W']  :
246            ptab [..., -1, 0:       ] = psgn * ptab [..., -2, -1::-1       ]
247               
248        if cd_type == 'U' :
249            ptab [..., -1, 0:-1     ] = psgn * ptab [..., -2, -2::-1       ]       
250            ptab [..., -1, -1       ] = psgn * ptab [..., -2, 0            ]
251             
252        if cd_type == 'V' :
253            ptab [..., -1, 0:       ] = psgn * ptab [..., -3, -1::-1       ]
254            ptab [..., -2, jpi//2:  ] = psgn * ptab [..., -2, jpi//2-1::-1 ]
255                             
256        if cd_type == 'F' :
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 ]
260
261    return ptab
262
263def lbc_mask (ptab, nperio=None, cd_type='T', sval=None) :
264    #
265    """
266    Mask fields on duplicated points
267    ptab      : Input array
268      rank 2 at least : ptab[...., lat, lon]
269    nperio    : Type of periodicity
270      1, 4, 6 : Cyclic on i dimension (generaly longitudes)
271      2       : Obsolete (was symmetric condition at southern boundary ?)
272      3, 4    : North fold T-point pivot (legacy ORCA2)
273      5, 6    : North fold F-point pivot (ORCA1, ORCA025, ORCA2 with new grid for paleo)
274    cd_type   : Grid specification : T, U, V or F
275   
276    See NEMO documentation for further details
277    """
278
279    jpi    = ptab.shape[-1]
280    nperio = __guessNperio__ (jpi, nperio)
281    if sval == None : sval = ptab.dtype.type (0)
282
283    #
284    #> East-West boundary conditions
285    # ------------------------------
286    if nperio in [1, 4, 6] :
287        # ... cyclic
288        ptab [...,:,  0] = sval
289        ptab [...,:, -1] = sval
290
291    #
292    #> South (in which nperio cases ?)
293    # --------------------------------
294    if nperio in [1, 3, 4, 5, 6] :
295        ptab [...,0,:] = sval
296       
297    #
298    #> North-South boundary conditions
299    # --------------------------------
300    if nperio in [3, 4] :  # North fold T-point pivot
301        if cd_type in [ 'T', 'W' ] : # T-, W-point
302            ptab [..., -1, 1:       ] = sval
303            ptab [..., -1, 0        ] = sval
304            ptab [..., -2, jpi//2:  ] = sval
305                     
306        if cd_type == 'U' :
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
316           
317        if cd_type == 'F' :
318            ptab [..., -2, 0:-1     ] = sval
319            ptab [..., -1, 0:-1     ] = sval
320            ptab [..., -1,  0       ] = sval
321            ptab [..., -1, -1       ] = sval
322     
323    if nperio in [5, 6] :            #  North fold F-point pivot
324        if cd_type in ['T', 'W']  :
325            ptab [..., -1, 0:       ] = sval
326               
327        if cd_type == 'U' :
328            ptab [..., -1, 0:-1     ] = sval       
329            ptab [..., -1, -1       ] = sval
330             
331        if cd_type == 'V' :
332            ptab [..., -1, 0:       ] = sval
333            ptab [..., -2, jpi//2:  ] = sval
334                             
335        if cd_type == 'F' :
336            ptab [..., -1, 0:-1     ] = sval
337            ptab [..., -1, -1       ] = sval
338            ptab [..., -2, jpi//2:-1] = sval
339
340    return ptab
341
342## ===========================================================================
343##
344##                               That's all folk's !!!
345##
346## ===========================================================================
Note: See TracBrowser for help on using the repository browser.