source: TOOLS/MOSAIX/nemo.py @ 6190

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

O.M. : Evolution on MOSAIX

  • Add licensing information
  • Update and beautify README.md
  • Add few model version to known list
  • Property svn:keywords set to Date Revision HeadURL Author Id
File size: 55.8 KB
Line 
1# -*- coding: utf-8 -*-
2## ===========================================================================
3##
4##  MOSAIX is under CeCILL_V2 licence. See "Licence_CeCILL_V2-en.txt"
5##  file for an english version of the licence and
6##  "Licence_CeCILL_V2-fr.txt" for a french version.
7##
8##  Permission is hereby granted, free of charge, to any person or
9##  organization obtaining a copy of the software and accompanying
10##  documentation covered by this license (the "Software") to use,
11##  reproduce, display, distribute, execute, and transmit the
12##  Software, and to prepare derivative works of the Software, and to
13##  permit third-parties to whom the Software is furnished to do so,
14##  all subject to the following:
15##
16##  Warning, to install, configure, run, use any of MOSAIX software or
17##  to read the associated documentation you'll need at least one (1)
18##  brain in a reasonably working order. Lack of this implement will
19##  void any warranties (either express or implied).  Authors assumes
20##  no responsability for errors, omissions, data loss, or any other
21##  consequences caused directly or indirectly by the usage of his
22##  software by incorrectly or partially configured
23##
24## ===========================================================================
25'''
26Utilities to plot NEMO ORCA fields
27Periodicity and other stuff
28
29olivier.marti@lsce.ipsl.fr
30'''
31
32## SVN information
33__Author__   = "$Author$"
34__Date__     = "$Date$"
35__Revision__ = "$Revision$"
36__Id__       = "$Id$"
37__HeadURL    = "$HeadURL$"
38
39import sys, numpy as np
40try    : import xarray as xr
41except : pass
42
43rpi = np.pi ; rad = np.deg2rad (1.0) ; dar = np.rad2deg (1.0)
44
45nperio_valid_range = [0, 1, 4, 4.2, 5, 6, 6.2]
46
47rday   = 24.*60.*60.     # Day length [s]
48rsiyea = 365.25 * rday * 2. * rpi / 6.283076 # Sideral year length [s]
49rsiday = rday / (1. + rday / rsiyea)
50raamo  =  12.        # Number of months in one year
51rjjhh  =  24.        # Number of hours in one day
52rhhmm  =  60.        # Number of minutes in one hour
53rmmss  =  60.        # Number of seconds in one minute
54omega  = 2. * rpi / rsiday # Earth rotation parameter [s-1]
55ra     = 6371229.    # Earth radius [m]
56grav   = 9.80665     # Gravity [m/s2]
57repsi  = np.finfo (1.0).eps
58
59xList = [ 'x', 'X', 'lon'   , 'longitude' ]
60yList = [ 'y', 'Y', 'lat'   , 'latitude'  ]
61zList = [ 'z', 'Z', 'depth' , ]
62tList = [ 't', 'T', 'time'  , ]
63
64## SVN information
65__Author__   = "$Author$"
66__Date__     = "$Date$"
67__Revision__ = "$Revision$"
68__Id__       = "$Id$"
69__HeadURL    = "$HeadURL$"
70
71def __math__ (tab, default=None) :
72    math = default
73    try    :
74        if type (tab) == xr.core.dataarray.DataArray : math = xr
75    except :
76        pass
77
78    try    :
79        if type (tab) == np.ndarray : math = np
80    except :
81        pass
82           
83    return math
84
85def __guessNperio__ (jpj, jpi, nperio=None) :
86    '''
87    Tries to guess the value of nperio (periodicity parameter. See NEMO documentation for details)
88   
89    Inputs
90    jpj    : number of latitudes
91    jpi    : number of longitudes
92    nperio : periodicity parameter
93    '''
94    if nperio == None :
95        ## Values for NEMO version < 4.2
96        if jpi ==  182 :
97            nperio = 4   # ORCA2. We choose legacy orca2.
98            Iperio = 1 ; Jperio = 0 ; NFold = 1 ; NFtype = 'T'
99        if jpi ==  362 : # eORCA1.
100            nperio = 6 
101            Iperio = 1 ; Jperio = 0 ; NFold = 1 ; NFtype = 'F'
102        if jpi == 1442 :  # ORCA025.
103            nperio = 6 
104            Iperio = 1 ; Jperio = 0 ; NFold = 1 ; NFtype = 'F'
105        if jpj ==  149 : # ORCA2. We choose legacy orca2.
106            nperio = 4
107            Iperio = 1 ; Jperio = 0 ; NFold = 1 ; NFtype = 'F'
108        if jpj ==  294 : # ORCA1
109            nperio = 6
110            Iperio = 1 ; Jperio = 0 ; NFold = 1 ; NFtype = 'F'
111        if jpj ==  332 : # eORCA1.
112            nperio = 6 
113            Iperio = 1 ; Jperio = 0 ; NFold = 1 ; NFtype = 'F'
114           
115        ## Values for NEMO version >= 4.2. No more halo points
116        if jpi ==  180 :
117            nperio = 4.2 # ORCA2. We choose legacy orca2.
118            Iperio = 1 ; Jperio = 0 ; NFold = 1 ; NFtype = 'F'
119        if jpi ==  360 : # eORCA1.
120            nperio = 6.2
121            Iperio = 1 ; Jperio = 0 ; NFold = 1 ; NFtype = 'F'
122        if jpi == 1440 : # ORCA025.
123            nperio = 6.2
124            Iperio = 1 ; Jperio = 0 ; NFold = 1 ; NFtype = 'F'
125
126        if jpj ==  148 : nperio = 4.2 # ORCA2. We choose legacy orca2.
127        if jpj ==  330 : nperio = 6.2 # eORCA1.
128
129        if nperio == None :
130            sys.exit ('in nemo module : nperio not found, and cannot by guessed')
131        else :
132            if nperio in nperio_valid_range :
133                print ('nperio set as {:} (deduced from jpi={:d})'.format (nperio, jpi))
134            else : 
135                sys.exit  ('nperio set as {:} (deduced from jpi={:d}) : nemo.py is not ready for this value'.format (nperio, jpi))
136           
137    return nperio
138
139def __guessPoint__ (ptab) :
140    '''
141    Tries to guess the grid point (periodicity parameter. See NEMO documentation for details)
142   
143    For array conforments with xgcm requirements
144
145    Inputs
146         ptab : xarray array
147
148    Credits : who is the original author ?
149    '''
150    gP = None
151    math = __math__ (ptab)
152    if math == xr :
153        if 'x_c' in ptab.dims and 'y_c' in ptab.dims                        : gP = 'T'
154        if 'x_f' in ptab.dims and 'y_c' in ptab.dims                        : gP = 'U'
155        if 'x_c' in ptab.dims and 'y_f' in ptab.dims                        : gP = 'V'
156        if 'x_f' in ptab.dims and 'y_f' in ptab.dims                        : gP = 'F'
157        if 'x_c' in ptab.dims and 'y_c' in ptab.dims and 'z_c' in ptab.dims : gP = 'T'
158        if 'x_c' in ptab.dims and 'y_c' in ptab.dims and 'z_f' in ptab.dims : gP = 'W'
159        if 'x_f' in ptab.dims and 'y_c' in ptab.dims and 'z_f' in ptab.dims : gP = 'U'
160        if 'x_c' in ptab.dims and 'y_f' in ptab.dims and 'z_f' in ptab.dims : gP = 'V'
161        if 'x_f' in ptab.dims and 'y_f' in ptab.dims and 'z_f' in ptab.dims : gP = 'F'
162             
163        if gP == None :
164            sys.exit ('in nemo module : cd_type not found, and cannot by guessed')
165        else :
166            print ('Grid set as', gP, 'deduced from dims ', ptab.dims)
167            return gP
168    else :
169         sys.exit ('in nemo module : cd_type not found, input is not an xarray data')
170         #return gP
171
172def __findAxis__ (tab, axis='z') :
173    '''
174    Find number and name of the requested axis
175    '''
176    math = __math__ (tab)
177    ix = None ; ax = None
178
179    if axis in xList :
180        axList = [ 'x', 'X',
181                   'lon', 'nav_lon', 'nav_lon_T', 'nav_lon_U', 'nav_lon_V', 'nav_lon_F', 'nav_lon_W',
182                   'x_grid_T', 'x_grid_U', 'x_grid_V', 'x_grid_F', 'x_grid_W',
183                   'glam', 'glamt', 'glamu', 'glamv', 'glamf', 'glamw' ]
184        unList = [ 'degrees_east' ]
185    if axis in yList :
186        axList = [ 'y', 'Y', 'lat',
187                   'nav_lat', 'nav_lat_T', 'nav_lat_U', 'nav_lat_V', 'nav_lat_F', 'nav_lat_W',
188                   'y_grid_T', 'y_grid_U', 'y_grid_V', 'y_grid_F', 'y_grid_W',
189                   'gphi', 'gphi', 'gphiu', 'gphiv', 'gphif', 'gphiw']
190        unList = [ 'degrees_north' ]
191    if axis in zList :
192        axList = [ 'z', 'Z',
193                   'depth', 'deptht', 'depthu', 'depthv', 'depthf', 'depthw',
194                   'olevel' ]
195        unList = [ 'm', 'meter' ]
196    if axis in tList :
197        axList = [ 't', 'T', 'time', 'time_counter' ]
198        unList = [ 'second', 'minute', 'hour', 'day', 'month' ]
199   
200    if math == xr :
201        for Name in axList :
202            try    :
203                ix = tab.dims.index (Name)
204                ax = Name
205            except : pass
206
207        for i, dim in enumerate (tab.dims) :
208            if 'units' in tab.coords[dim].attrs.keys() :
209                for name in unList :
210                    if name in tab.coords[dim].attrs['units'] :
211                        ix = i
212                        ax = dim
213    else :
214        if axis in xList : ix=-1
215        if axis in yList :
216            if len(tab.shape) >= 2 : ix=-2
217        if axis in zList :
218            if len(tab.shape) >= 3 : ix=-3
219        if axis in tList :
220            if len(tab.shape) >=3  : ix=-3
221            if len(tab.shape) >=4  : ix=-4
222       
223    return ix, ax
224               
225def fixed_lon (lon, center_lon=0.0) :
226    '''
227    Returns corrected longitudes for nicer plots
228
229    lon        : longitudes of the grid. At least 2D.
230    center_lon : center longitude. Default=0.
231
232    Designed by Phil Pelson. See https://gist.github.com/pelson/79cf31ef324774c97ae7
233    '''
234    math = __math__ (lon)
235   
236    fixed_lon = lon.copy ()
237       
238    fixed_lon = math.where (fixed_lon > center_lon+180., fixed_lon-360.0, fixed_lon)
239    fixed_lon = math.where (fixed_lon < center_lon-180., fixed_lon+360.0, fixed_lon)
240   
241    for i, start in enumerate (np.argmax (np.abs (np.diff (fixed_lon, axis=-1)) > 180., axis=-1)) :
242        fixed_lon [..., i, start+1:] += 360
243
244    # Special case for eORCA025
245    if fixed_lon.shape [-1] == 1442 : fixed_lon [..., -2, :] = fixed_lon [..., -3, :]
246    if fixed_lon.shape [-1] == 1440 : fixed_lon [..., -1, :] = fixed_lon [..., -2, :]
247
248    if fixed_lon.min () > center_lon : fixed_lon = fixed_lon-360.0
249               
250    return fixed_lon
251
252def jeq (lat) :
253    '''
254    Returns j index of equator in the grid
255   
256    lat : latitudes of the grid. At least 2D.
257    '''
258    math = __math__ (lat)
259    ix, ax = __findAxis__ (lat, 'x')
260    iy, ay = __findAxis__ (lat, 'y')
261
262    if math == xr :
263        jeq = int ( np.mean ( np.argmin (np.abs (np.float64 (lat)), axis=iy) ) )
264    else : 
265        jeq = np.argmin (np.abs (np.float64 (lat[...,:, 0])))
266    return jeq
267
268def lon1D (lon, lat=None) :
269    '''
270    Returns 1D longitude for simple plots.
271   
272    lon : longitudes of the grid
273    lat (optionnal) : latitudes of the grid
274    '''
275    if np.max (lat) != None :
276        je    = jeq (lat)
277        lon1D = lon [..., je, :]
278    else :
279        jpj, jpi = lon.shape [-2:]
280        lon1D    = lon [..., jpj//3, :]
281
282    start = np.argmax (np.abs (np.diff (lon1D, axis=-1)) > 180.0, axis=-1)
283    lon1D [..., start+1:] += 360
284       
285    return lon1D
286
287def latreg (lat, diff=0.1) :
288    '''
289    Returns maximum j index where gridlines are along latitude in the northern hemisphere
290   
291    lat : latitudes of the grid (2D)
292    diff [optional] : tolerance
293    '''
294    math = __math__ (lat)
295    if diff == None :
296        dy   = np.float64 (np.mean (np.abs (lat - np.roll (lat,shift=1,axis=-2, roll_coords=False))))
297        diff = dy/100.
298   
299    je     = jeq (lat)
300    jreg   = np.where (lat[...,je:,:].max(axis=-1) - lat[...,je:,:].min(axis=-1)< diff)[-1][-1] + je
301    latreg = np.float64 (lat[...,jreg,:].mean(axis=-1))
302    JREG   = jreg
303
304    return jreg, latreg
305
306def lat1D (lat) :
307    '''
308    Returns 1D latitudes for zonal means and simple plots.
309
310    lat : latitudes of the grid (2D)
311    '''
312    math = __math__ (lat)
313    jpj, jpi = lat.shape[-2:]
314
315    dy     = np.float64 (np.mean (np.abs (lat - np.roll (lat, shift=1,axis=-2))))
316    je     = jeq (lat)
317    lat_eq = np.float64 (lat[...,je,:].mean(axis=-1))
318     
319    jreg, lat_reg = latreg (lat)
320    lat_ave = np.mean (lat, axis=-1)
321
322    if (np.abs (lat_eq) < dy/100.) : # T, U or W grid
323        dys    = (90.-lat_reg) / (jpj-jreg-1)*0.5
324        yrange = 90.-dys-lat_reg
325    else                           :  # V or F grid
326        yrange = 90.    -lat_reg
327       
328    lat1D = math.where (lat_ave<lat_reg, lat_ave, lat_reg + yrange * (np.arange(jpj)-jreg)/(jpj-jreg-1))   
329       
330    if math == xr : lat1D.attrs = lat.attrs
331
332    return lat1D
333
334def latlon1D (lat, lon) :
335    '''
336    Returns simple latitude and longitude (1D) for simple plots.
337
338    lat, lon : latitudes and longitudes of the grid (2D)
339    '''
340    return lat1D (lat),  lon1D (lon, lat)
341
342def mask_lonlat (ptab, x0, x1, y0, y1, lon, lat, sval=np.nan) :
343    math = __math__ (ptab)
344    try :
345        lon = lon.copy().to_masked_array()
346        lat = lat.copy().to_masked_array()
347    except : pass
348           
349    mask = np.logical_and (np.logical_and(lat>y0, lat<y1), 
350            np.logical_or (np.logical_or (np.logical_and(lon>x0, lon<x1), np.logical_and(lon+360>x0, lon+360<x1)),
351                                      np.logical_and(lon-360>x0, lon-360<x1)))
352    tab = math.where (mask, ptab, np.nan)
353   
354    return tab
355       
356def extend (tab, Lon=False, jplus=25, jpi=None, nperio=4) :
357    '''
358    Returns extended field eastward to have better plots, and box average crossing the boundary
359    Works only for xarray and numpy data (?)
360
361    tab : field to extend.
362    Lon : (optional, default=False) : if True, add 360 in the extended parts of the field
363    jpi : normal longitude dimension of the field. exrtend does nothing it the actual
364        size of the field != jpi (avoid to extend several times)
365    jplus (optional, default=25) : number of points added on the east side of the field
366   
367    '''
368    math = __math__ (tab)
369   
370    if tab.shape[-1] == 1 : extend = tab
371
372    else :
373        if jpi == None : jpi = tab.shape[-1]
374
375        if Lon : xplus = -360.0
376        else   : xplus =    0.0
377
378        if tab.shape[-1] > jpi :
379            extend = tab
380        else :
381            if nperio == 0 or nperio == 4.2 :
382                istart = 0 ; le=jpi+1 ; la=0
383            if nperio == 1 :
384                istart = 0 ; le=jpi+1 ; la=0
385            if nperio == 4 or nperio == 6 : # OPA case with two halo points for periodicity
386                istart = 1 ; le=jpi-2 ; la=1  # Perfect, except at the pole that should be masked by lbc_plot
387           
388            if math == xr :
389                extend = np.concatenate ((tab.values[..., istart   :istart+le+1    ] + xplus,
390                                          tab.values[..., istart+la:istart+la+jplus]         ), axis=-1)
391                lon    = tab.dims[-1]
392                new_coords = []
393                for coord in tab.dims :
394                    if coord == lon : new_coords.append ( np.arange( extend.shape[-1]))
395                    else            : new_coords.append ( tab.coords[coord].values)
396                extend = xr.DataArray ( extend, dims=tab.dims, coords=new_coords )
397            else : 
398                extend = np.concatenate ((tab [..., istart   :istart+le+1    ] + xplus,
399                                          tab [..., istart+la:istart+la+jplus]          ), axis=-1)
400    return extend
401
402def orca2reg (ff, lat_name='nav_lat', lon_name='nav_lon', y_name='y', x_name='x') :
403    '''
404    Assign ORCA dataset on a regular grid.
405    For use in the tropical region.
406
407    Inputs :
408      ff : xarray dataset
409      lat_name, lon_name : name of latitude and longitude 2D field in ff
410      y_name, x_name     : namex of dimensions in ff
411
412    Returns : xarray dataset with rectangular grid. Incorrect above 20°N
413    '''
414    # Compute 1D longitude and latitude
415    (lat, lon) = latlon1D (ff[lat_name], ff[lon_name])
416    # Assign lon and lat as dimensions of the dataset
417    if y_name in ff.dims : 
418        lat = xr.DataArray (lat, coords=[lat,], dims=['lat',])     
419        ff  = ff.rename_dims ({y_name: "lat",}).assign_coords (lat=lat)
420    if x_name in ff.dims :
421        lon = xr.DataArray (lon, coords=[lon,], dims=['lon',])
422        ff  = ff.rename_dims ({x_name: "lon",}).assign_coords (lon=lon)
423    # Force dimensions to be in the right order
424    coord_order = ['lat', 'lon']
425    for dim in [ 'depthw', 'depthv', 'depthu', 'deptht', 'depth', 'z',
426                 'time_counter', 'time', 'tbnds', 
427                 'bnds', 'axis_nbounds', 'two2', 'two1', 'two', 'four',] :
428        if dim in ff.dims : coord_order.insert (0, dim)
429       
430    ff = ff.transpose (*coord_order)
431    return ff
432
433def lbc_init (ptab, nperio=None) :
434    '''
435    Prepare for all lbc calls
436   
437    Set periodicity on input field
438    nperio    : Type of periodicity
439      0       : No periodicity
440      1, 4, 6 : Cyclic on i dimension (generaly longitudes) with 2 points halo
441      2       : Obsolete (was symmetric condition at southern boundary ?)
442      3, 4    : North fold T-point pivot (legacy ORCA2)
443      5, 6    : North fold F-point pivot (ORCA1, ORCA025, ORCA2 with new grid for paleo)
444    cd_type   : Grid specification : T, U, V or F
445
446    See NEMO documentation for further details
447    '''
448    jpj, jpi = ptab.shape[-2:]
449    if nperio == None : nperio = __guessNperio__ (jpj, jpi, nperio)
450   
451    if nperio not in nperio_valid_range :
452        print ('nperio=', nperio, ' is not in the valid range', nperio_valid_range)
453        sys.exit ()
454
455    return jpj, jpi, nperio
456       
457       
458def lbc (ptab, nperio=None, cd_type='T', psgn=1.0, nemo_4U_bug=False) :
459    '''
460    Set periodicity on input field
461    ptab      : Input array (works for rank 2 at least : ptab[...., lat, lon])
462    nperio    : Type of periodicity
463    cd_type   : Grid specification : T, U, V or F
464    psgn      : For change of sign for vector components (1 for scalars, -1 for vector components)
465   
466    See NEMO documentation for further details
467    '''
468    jpj, jpi, nperio = lbc_init (ptab, nperio)
469    psgn   = ptab.dtype.type (psgn)
470    math = __math__ (ptab)
471   
472    if math == xr : ztab = ptab.values.copy ()
473    else          : ztab = ptab.copy ()
474    #
475    #> East-West boundary conditions
476    # ------------------------------
477    if nperio in [1, 4, 6] :
478        # ... cyclic
479        ztab [..., :,  0] = ztab [..., :, -2]
480        ztab [..., :, -1] = ztab [..., :,  1]
481    #
482    #> North-South boundary conditions
483    # --------------------------------
484    if nperio in [3, 4] :  # North fold T-point pivot
485        if cd_type in [ 'T', 'W' ] : # T-, W-point
486            ztab [..., -1, 1:       ] = psgn * ztab [..., -3, -1:0:-1      ]
487            ztab [..., -1, 0        ] = psgn * ztab [..., -3, 2            ]
488            ztab [..., -2, jpi//2:  ] = psgn * ztab [..., -2, jpi//2:0:-1  ]
489               
490        if cd_type == 'U' :
491            ztab [..., -1, 0:-1     ] = psgn * ztab [..., -3, -1:0:-1      ]       
492            ztab [..., -1,  0       ] = psgn * ztab [..., -3,  1           ]
493            ztab [..., -1, -1       ] = psgn * ztab [..., -3, -2           ]
494               
495            if nemo_4U_bug :
496                ztab [..., -2, jpi//2+1:-1] = psgn * ztab [..., -2, jpi//2-2:0:-1]
497                ztab [..., -2, jpi//2-1   ] = psgn * ztab [..., -2, jpi//2       ]
498            else :
499                ztab [..., -2, jpi//2-1:-1] = psgn * ztab [..., -2, jpi//2:0:-1]
500               
501        if cd_type == 'V' : 
502            ztab [..., -2, 1:       ] = psgn * ztab [..., -3, jpi-1:0:-1   ]
503            ztab [..., -1, 1:       ] = psgn * ztab [..., -4, -1:0:-1      ]   
504            ztab [..., -1, 0        ] = psgn * ztab [..., -4, 2            ]
505               
506        if cd_type == 'F' :
507            ztab [..., -2, 0:-1     ] = psgn * ztab [..., -3, -1:0:-1      ]
508            ztab [..., -1, 0:-1     ] = psgn * ztab [..., -4, -1:0:-1      ]
509            ztab [..., -1,  0       ] = psgn * ztab [..., -4,  1           ]
510            ztab [..., -1, -1       ] = psgn * ztab [..., -4, -2           ]
511
512    if nperio in [4.2] :  # North fold T-point pivot
513        if cd_type in [ 'T', 'W' ] : # T-, W-point
514            ztab [..., -1, jpi//2:  ] = psgn * ztab [..., -1, jpi//2:0:-1  ]
515               
516        if cd_type == 'U' :
517            ztab [..., -1, jpi//2-1:-1] = psgn * ztab [..., -1, jpi//2:0:-1]
518               
519        if cd_type == 'V' : 
520            ztab [..., -1, 1:       ] = psgn * ztab [..., -2, jpi-1:0:-1   ]
521               
522        if cd_type == 'F' :
523            ztab [..., -1, 0:-1     ] = psgn * ztab [..., -2, -1:0:-1      ]
524
525    if nperio in [5, 6] :            #  North fold F-point pivot 
526        if cd_type in ['T', 'W']  :
527            ztab [..., -1, 0:       ] = psgn * ztab [..., -2, -1::-1       ]
528               
529        if cd_type == 'U' :
530            ztab [..., -1, 0:-1     ] = psgn * ztab [..., -2, -2::-1       ]       
531            ztab [..., -1, -1       ] = psgn * ztab [..., -2, 0            ] # Bug ?
532               
533        if cd_type == 'V' :
534            ztab [..., -1, 0:       ] = psgn * ztab [..., -3, -1::-1       ]
535            ztab [..., -2, jpi//2:  ] = psgn * ztab [..., -2, jpi//2-1::-1 ]
536               
537        if cd_type == 'F' :
538            ztab [..., -1, 0:-1     ] = psgn * ztab [..., -3, -2::-1       ]
539            ztab [..., -1, -1       ] = psgn * ztab [..., -3, 0            ]
540            ztab [..., -2, jpi//2:-1] = psgn * ztab [..., -2, jpi//2-2::-1 ]
541
542    #
543    #> East-West boundary conditions
544    # ------------------------------
545    if nperio in [1, 4, 6] :
546        # ... cyclic
547        ztab [..., :,  0] = ztab [..., :, -2]
548        ztab [..., :, -1] = ztab [..., :,  1]
549
550    if math == xr :
551        ztab = xr.DataArray ( ztab, dims=ptab.dims, coords=ptab.coords)
552       
553    return ztab
554
555def lbc_mask (ptab, nperio=None, cd_type='T', sval=np.nan) :
556    #
557    '''
558    Mask fields on duplicated points
559    ptab      : Input array. Rank 2 at least : ptab [...., lat, lon]
560    nperio    : Type of periodicity
561    cd_type   : Grid specification : T, U, V or F
562   
563    See NEMO documentation for further details
564    '''
565    jpj, jpi, nperio = lbc_init (ptab, nperio)
566    ztab = ptab.copy ()
567
568    #
569    #> East-West boundary conditions
570    # ------------------------------
571    if nperio in [1, 4, 6] :
572        # ... cyclic
573        ztab [..., :,  0] = sval
574        ztab [..., :, -1] = sval
575
576    #
577    #> South (in which nperio cases ?)
578    # --------------------------------
579    if nperio in [1, 3, 4, 5, 6] :
580        ztab [..., 0, :] = sval
581       
582    #
583    #> North-South boundary conditions
584    # --------------------------------
585    if nperio in [3, 4] :  # North fold T-point pivot
586        if cd_type in [ 'T', 'W' ] : # T-, W-point
587            ztab [..., -1,  :         ] = sval
588            ztab [..., -2, :jpi//2  ] = sval
589               
590        if cd_type == 'U' :
591            ztab [..., -1,  :         ] = sval 
592            ztab [..., -2, jpi//2+1:  ] = sval
593               
594        if cd_type == 'V' :
595            ztab [..., -2, :       ] = sval
596            ztab [..., -1, :       ] = sval   
597               
598        if cd_type == 'F' :
599            ztab [..., -2, :       ] = sval
600            ztab [..., -1, :       ] = sval
601
602    if nperio in [4.2] :  # North fold T-point pivot
603        if cd_type in [ 'T', 'W' ] : # T-, W-point
604            ztab [..., -1, jpi//2  :  ] = sval
605               
606        if cd_type == 'U' :
607            ztab [..., -1, jpi//2-1:-1] = sval
608               
609        if cd_type == 'V' : 
610            ztab [..., -1, 1:       ] = sval
611               
612        if cd_type == 'F' :
613            ztab [..., -1, 0:-1     ] = sval
614   
615    if nperio in [5, 6] :            #  North fold F-point pivot
616        if cd_type in ['T', 'W']  :
617            ztab [..., -1, 0:       ] = sval
618               
619        if cd_type == 'U' :
620            ztab [..., -1, 0:-1     ] = sval       
621            ztab [..., -1, -1       ] = sval
622             
623        if cd_type == 'V' :
624            ztab [..., -1, 0:       ] = sval
625            ztab [..., -2, jpi//2:  ] = sval
626                             
627        if cd_type == 'F' :
628            ztab [..., -1, 0:-1       ] = sval
629            ztab [..., -1, -1         ] = sval
630            ztab [..., -2, jpi//2+1:-1] = sval
631
632    return ztab
633
634def lbc_plot (ptab, nperio=None, cd_type='T', psgn=1.0, sval=np.nan) :
635    '''
636    Set periodicity on input field, adapted for plotting for any cartopy projection
637    ptab      : Input array. Rank 2 at least : ptab[...., lat, lon]
638    nperio    : Type of periodicity
639    cd_type   : Grid specification : T, U, V or F
640    psgn      : For change of sign for vector components (1 for scalars, -1 for vector components)
641   
642    See NEMO documentation for further details
643    '''
644
645    jpj, jpi, nperio = lbc_init (ptab, nperio)
646    psgn   = ptab.dtype.type (psgn)
647    ztab   = ptab.copy ()
648    #
649    #> East-West boundary conditions
650    # ------------------------------
651    if nperio in [1, 4, 6] :
652        # ... cyclic
653        ztab [..., :,  0] = ztab [..., :, -2]
654        ztab [..., :, -1] = ztab [..., :,  1]
655
656    #> Masks south
657    # ------------
658    if nperio in [4, 6] : ztab [..., 0, : ] = sval
659       
660    #
661    #> North-South boundary conditions
662    # --------------------------------
663    if nperio in [3, 4] :  # North fold T-point pivot
664        if cd_type in [ 'T', 'W' ] : # T-, W-point
665            ztab [..., -1,  :      ] = sval
666            #ztab [..., -2, jpi//2: ] = sval
667            ztab [..., -2, :jpi//2 ] = sval # Give better plots than above
668        if cd_type == 'U' :
669            ztab [..., -1, : ] = sval
670
671        if cd_type == 'V' : 
672            ztab [..., -2, : ] = sval
673            ztab [..., -1, : ] = sval
674           
675        if cd_type == 'F' :
676            ztab [..., -2, : ] = sval
677            ztab [..., -1, : ] = sval
678
679    if nperio in [4.2] :  # North fold T-point pivot
680        if cd_type in [ 'T', 'W' ] : # T-, W-point
681            ztab [..., -1, jpi//2:  ] = sval
682               
683        if cd_type == 'U' :
684            ztab [..., -1, jpi//2-1:-1] = sval
685               
686        if cd_type == 'V' : 
687            ztab [..., -1, 1:       ] = sval
688               
689        if cd_type == 'F' :
690            ztab [..., -1, 0:-1     ] = sval
691     
692    if nperio in [5, 6] :            #  North fold F-point pivot 
693        if cd_type in ['T', 'W']  :
694            ztab [..., -1, : ] = sval
695               
696        if cd_type == 'U' :
697            ztab [..., -1, : ] = sval     
698             
699        if cd_type == 'V' :
700            ztab [..., -1, :        ] = sval
701            ztab [..., -2, jpi//2:  ] = sval
702                             
703        if cd_type == 'F' :
704            ztab [..., -1, :          ] = sval
705            ztab [..., -2, jpi//2+1:-1] = sval
706
707    return ztab
708
709def lbc_add (ptab, nperio=None, cd_type =None) :
710    '''
711    Handle NEMO domain changes between NEMO 4.0 to NEMO 4.2
712      Peridodicity halo has been removes
713    This routine adds the halos if needed
714
715    ptab      : Input array (works
716      rank 2 at least : ptab[...., lat, lon]
717    nperio    : Type of periodicity
718 
719    See NEMO documentation for further details
720    '''
721
722    jpj, jpi, nperio = lbc_init (ptab, nperio)
723
724    t_shape = np.array (ptab.shape)
725
726    if nperio == 4.2 or nperio == 6.2 :
727        math = __math__ (ptab) 
728     
729        ext_shape = t_shape
730        ext_shape[-1] = ext_shape[-1] + 2
731        ext_shape[-2] = ext_shape[-2] + 1
732
733        if math == xr : ptab_ext = xr.DataArray (np.empty (ext_shape), dims=ptab.dims) 
734        else          : ptab_ext =               np.empty (ext_shape)
735           
736        ptab_ext[..., :-1, 1:-1] = ptab
737       
738        if nperio == 4.2 : ptab_ext = lbc (ptab_ext, 4, cd_type)
739        if nperio == 6.2 : ptab_ext = lbc (ptab_ext, 6, cd_type)
740
741        if math == xr :
742            for attr in ptab.attrs :
743                ptab_ext.attrs [attr] = ptab.attrs [attr]
744
745    else : ptab_ext = ptab
746       
747    return ptab
748
749def lbc_del (ptab, nperio=None) :
750    '''
751    Handle NEMO domain changes between NEMO 4.0 to NEMO 4.2
752      Peridodicity halo has been removes
753    This routine removes the halos if needed
754
755    ptab      : Input array (works
756      rank 2 at least : ptab[...., lat, lon]
757    nperio    : Type of periodicity
758 
759    See NEMO documentation for further details
760    '''
761
762    jpj, jpi, nperio = lbc_init (ptab, nperio)
763
764    if nperio == 4 or nperio == 6 :
765        return ptab[..., :-1, 1:-1]
766    else :
767        return ptab
768
769def lbc_index (jj, ii, jpj, jpi, nperio=None, cd_type='T') :
770    '''
771    For indexes of a NEMO point, give the corresponding point inside the util domain
772    jj, ii    : indexes
773    jpi, jpi  : size of domain
774    nperio    : type of periodicity
775    cd_type   : grid specification : T, U, V or F
776   
777    See NEMO documentation for further details
778    '''
779
780    if nperio == None : nperio = __guessNperio__ (jpj, jpi, nperio)
781   
782    ## For the sake of simplicity, switch to the convention of original lbc Fortran routine from NEMO
783    ## : starts indexes at 1
784    jy = jj + 1 ; ix = ii + 1
785
786    math = __math__ (jj)
787    if math == None : math=np
788
789    #
790    #> East-West boundary conditions
791    # ------------------------------
792    if nperio in [1, 4, 6] :
793        #... cyclic
794        ix = math.where (ix==jpi, 2   , ix)
795        ix = math.where (ix== 1 ,jpi-1, ix)
796
797    #
798    def modIJ (cond, jy_new, ix_new) :
799        jy_r = math.where (cond, jy_new, jy)
800        ix_r = math.where (cond, ix_new, ix)
801        return jy_r, ix_r
802    #
803    #> North-South boundary conditions
804    # --------------------------------
805    if nperio in [ 3 , 4 ]  :
806        if cd_type in  [ 'T' , 'W' ] :
807            (jy, ix) = modIJ (np.logical_and (jy==jpj  , ix>=2       ), jpj-2, jpi-ix+2)
808            (jy, ix) = modIJ (np.logical_and (jy==jpj  , ix==1       ), jpj-1, 3       )   
809            (jy, ix) = modIJ (np.logical_and (jy==jpj-1, ix>=jpi//2+1), jy   , jpi-ix+2) 
810
811        if cd_type in [ 'U' ] :
812            (jy, ix) = modIJ (np.logical_and (jy==jpj  , np.logical_and (ix>=1, ix <= jpi-1)   ), jy   , jpi-ix+1)
813            (jy, ix) = modIJ (np.logical_and (jy==jpj  , ix==1  )                               , jpj-2, 2       )
814            (jy, ix) = modIJ (np.logical_and (jy==jpj  , ix==jpi)                               , jpj-2, jpi-1   )
815            (jy, ix) = modIJ (np.logical_and (jy==jpj-1, np.logical_and (ix>=jpi//2, ix<=jpi-1)), jy   , jpi-ix+1)
816         
817        if cd_type in [ 'V' ] :
818            (jy, ix) = modIJ (np.logical_and (jy==jpj-1, ix>=2  ), jpj-2, jpi-ix+2)
819            (jy, ix) = modIJ (np.logical_and (jy==jpj  , ix>=2  ), jpj-3, jpi-ix+2)
820            (jy, ix) = modIJ (np.logical_and (jy==jpj  , ix==1  ), jpj-3,  3      )
821           
822        if cd_type in [ 'F' ] :
823            (jy, ix) = modIJ (np.logical_and (jy==jpj-1, ix<=jpi-1), jpj-2, jpi-ix+1)
824            (jy, ix) = modIJ (np.logical_and (jy==jpj  , ix<=jpi-1), jpj-3, jpi-ix+1)
825            (jy, ix) = modIJ (np.logical_and (jy==jpj  , ix==1    ), jpj-3, 2       )
826            (jy, ix) = modIJ (np.logical_and (jy==jpj  , ix==jpi  ), jpj-3, jpi-1   )
827
828    if nperio in [ 5 , 6 ] :
829        if cd_type in [ 'T' , 'W' ] :                        # T-, W-point
830             (jy, ix) = modIJ (jy==jpj, jpj-1, jpi-ix+1)
831 
832        if cd_type in [ 'U' ] :                              # U-point
833            (jy, ix) = modIJ (np.logical_and (jy==jpj  , ix<=jpi-1   ), jpj-1, jpi-ix  )
834            (jy, ix) = modIJ (np.logical_and (jy==jpj  , ix==jpi     ), jpi-1, 1       )
835           
836        if cd_type in [ 'V' ] :    # V-point
837            (jy, ix) = modIJ (jy==jpj                                 , jy   , jpi-ix+1)
838            (jy, ix) = modIJ (np.logical_and (jy==jpj-1, ix>=jpi//2+1), jy   , jpi-ix+1)
839           
840        if cd_type in [ 'F' ] :                              # F-point
841            (jy, ix) = modIJ (np.logical_and (jy==jpj  , ix<=jpi-1   ), jpj-2, jpi-ix  )
842            (jy, ix) = modIJ (np.logical_and (ix==jpj  , ix==jpi     ), jpj-2, 1       )
843            (jy, ix) = modIJ (np.logical_and (jy==jpj-1, ix>=jpi//2+1), jy   , jpi-ix  )
844
845    ## Restore convention to Python/C : indexes start at 0
846    jy += -1 ; ix += -1
847
848    if isinstance (jj, int) : jy = jy.item ()
849    if isinstance (ii, int) : ix = ix.item ()
850
851    return jy, ix
852   
853def geo2en (pxx, pyy, pzz, glam, gphi) : 
854    '''
855    Change vector from geocentric to east/north
856
857    Inputs :
858        pxx, pyy, pzz : components on the geocentric system
859        glam, gphi : longitude and latitude of the points
860    '''
861
862    gsinlon = np.sin (rad * glam)
863    gcoslon = np.cos (rad * glam)
864    gsinlat = np.sin (rad * gphi)
865    gcoslat = np.cos (rad * gphi)
866         
867    pte = - pxx * gsinlon            + pyy * gcoslon
868    ptn = - pxx * gcoslon * gsinlat  - pyy * gsinlon * gsinlat + pzz * gcoslat
869
870    return pte, ptn
871
872def en2geo (pte, ptn, glam, gphi) :
873    '''
874    Change vector from east/north to geocentric
875
876    Inputs :
877        pte, ptn   : eastward/northward components
878        glam, gphi : longitude and latitude of the points
879    '''
880   
881    gsinlon = np.sin (rad * glam)
882    gcoslon = np.cos (rad * glam)
883    gsinlat = np.sin (rad * gphi)
884    gcoslat = np.cos (rad * gphi)
885
886    pxx = - pte * gsinlon - ptn * gcoslon * gsinlat
887    pyy =   pte * gcoslon - ptn * gsinlon * gsinlat
888    pzz =   ptn * gcoslat
889   
890    return pxx, pyy, pzz
891
892def findJI (lat_data, lon_data, lat_grid, lon_grid, mask=1.0, verbose=False) :
893    '''
894    Description: seeks J,I indices of the grid point which is the closest of a given point
895    Usage: go FindJI  <data latitude> <data longitude> <grid latitudes> <grid longitudes> [mask]
896    <longitude fields> <latitude field> are 2D fields on J/I (Y/X) dimensions
897    mask : if given, seek only non masked grid points (i.e with mask=1)
898   
899    Example : findIJ (40, -20, nav_lat, nav_lon, mask=1.0)
900
901    Note : all longitudes and latitudes in degrees
902       
903    Note : may work with 1D lon/lat (?)
904    '''
905    # Get grid dimensions
906    if len (lon_grid.shape) == 2 : (jpj, jpi) = lon_grid.shape
907    else                         : jpj = len(lat_grid) ; jpi=len(lon_grid)
908
909    math = __math__ (lat_grid)
910       
911    # Compute distance from the point to all grid points (in radian)
912    arg      = np.sin (rad*lat_data) * np.sin (rad*lat_grid) \
913             + np.cos (rad*lat_data) * np.cos (rad*lat_grid) * np.cos(rad*(lon_data-lon_grid))
914    distance = np.arccos (arg) + 4.0*rpi*(1.0-mask) # Send masked points to 'infinite'
915
916    # Truncates to alleviate some precision problem with some grids
917    prec = int (1E7)
918    distance = (distance*prec).astype(int) / prec
919
920    # Compute minimum of distance, and index of minimum
921    #
922    distance_min = distance.min    ()
923    jimin        = int (distance.argmin ())
924   
925    # Compute 2D indices
926    jmin = jimin // jpi ; imin = jimin - jmin*jpi
927
928    # Compute distance achieved
929    mindist = distance[jmin, imin]
930   
931    # Compute azimuth
932    dlon = lon_data-lon_grid[jmin,imin]
933    arg  = np.sin (rad*dlon) /  (np.cos(rad*lat_data)*np.tan(rad*lat_grid[jmin,imin]) - np.sin(rad*lat_data)*np.cos(rad*dlon))
934    azimuth = dar*np.arctan (arg)
935   
936    # Result
937    if verbose : 
938        print ('I={:d} J={:d} - Data:{:5.1f}°N {:5.1f}°E - Grid:{:4.1f}°N {:4.1f}°E - Dist: {:6.1f}km {:5.2f}° - Azimuth: {:3.2f}rad - {:5.1f}°'
939            .format (imin, jmin, lat_data, lon_data, lat_grid[jmin,imin], lon_grid[jmin,imin], ra*distance[jmin,imin], dar*distance[jmin,imin], rad*azimuth, azimuth))
940
941    return jmin, imin
942
943def clo_lon (lon, lon0) :
944    '''Choose closest to lon0 longitude, adding or substacting 360° if needed'''
945    math = __math__ (lon, np)
946       
947    clo_lon = lon
948    clo_lon = math.where (clo_lon > lon0 + 180., clo_lon-360., clo_lon)
949    clo_lon = math.where (clo_lon < lon0 - 180., clo_lon+360., clo_lon)
950    clo_lon = math.where (clo_lon > lon0 + 180., clo_lon-360., clo_lon)
951    clo_lon = math.where (clo_lon < lon0 - 180., clo_lon+360., clo_lon)
952    if clo_lon.shape == () : clo_lon = clo_lon.item ()
953    return clo_lon
954
955def angle_full (glamt, gphit, glamu, gphiu, glamv, gphiv, glamf, gphif, nperio=None) :
956    '''Compute sinus and cosinus of model line direction with respect to east'''
957    math = __math__ (glamt)
958   
959    # north pole direction & modulous (at T-point)
960    zxnpt = 0. - 2.0 * np.cos (rad*glamt) * np.tan (rpi/4.0 - rad*gphit/2.0)
961    zynpt = 0. - 2.0 * np.sin (rad*glamt) * np.tan (rpi/4.0 - rad*gphit/2.0)
962    znnpt = zxnpt*zxnpt + zynpt*zynpt
963   
964    # north pole direction & modulous (at U-point)
965    zxnpu = 0. - 2.0 * np.cos (rad*glamu) * np.tan (rpi/4.0 - rad*gphiu/2.0)
966    zynpu = 0. - 2.0 * np.sin (rad*glamu) * np.tan (rpi/4.0 - rad*gphiu/2.0)
967    znnpu = zxnpu*zxnpu + zynpu*zynpu
968   
969    # north pole direction & modulous (at V-point)
970    zxnpv = 0. - 2.0 * np.cos (rad*glamv) * np.tan (rpi/4.0 - rad*gphiv/2.0)
971    zynpv = 0. - 2.0 * np.sin (rad*glamv) * np.tan (rpi/4.0 - rad*gphiv/2.0)
972    znnpv = zxnpv*zxnpv + zynpv*zynpv
973
974    # north pole direction & modulous (at F-point)
975    zxnpf = 0. - 2.0 * np.cos( rad*glamf ) * np.tan ( rpi/4. - rad*gphif/2. )
976    zynpf = 0. - 2.0 * np.sin( rad*glamf ) * np.tan ( rpi/4. - rad*gphif/2. )
977    znnpf = zxnpf*zxnpf + zynpf*zynpf
978
979    # j-direction: v-point segment direction (around T-point)
980    zlam = glamv 
981    zphi = gphiv
982    zlan = np.roll ( glamv, axis=-2, shift=1)  # glamv (ji,jj-1)
983    zphh = np.roll ( gphiv, axis=-2, shift=1)  # gphiv (ji,jj-1)
984    zxvvt =  2.0 * np.cos ( rad*zlam ) * np.tan ( rpi/4. - rad*zphi/2. )   \
985          -  2.0 * np.cos ( rad*zlan ) * np.tan ( rpi/4. - rad*zphh/2. )
986    zyvvt =  2.0 * np.sin ( rad*zlam ) * np.tan ( rpi/4. - rad*zphi/2. )   \
987          -  2.0 * np.sin ( rad*zlan ) * np.tan ( rpi/4. - rad*zphh/2. )
988    znvvt = np.sqrt ( znnpt * ( zxvvt*zxvvt + zyvvt*zyvvt )  )
989
990    # j-direction: f-point segment direction (around u-point)
991    zlam = glamf
992    zphi = gphif
993    zlan = np.roll (glamf, axis=-2, shift=1) # glamf (ji,jj-1)
994    zphh = np.roll (gphif, axis=-2, shift=1) # gphif (ji,jj-1)
995    zxffu =  2.0 * np.cos ( rad*zlam ) * np.tan ( rpi/4. - rad*zphi/2. )   \
996          -  2.0 * np.cos ( rad*zlan ) * np.tan ( rpi/4. - rad*zphh/2. )
997    zyffu =  2.0 * np.sin ( rad*zlam ) * np.tan ( rpi/4. - rad*zphi/2. )   \
998          -  2.0 * np.sin ( rad*zlan ) * np.tan ( rpi/4. - rad*zphh/2. )
999    znffu = np.sqrt ( znnpu * ( zxffu*zxffu + zyffu*zyffu )  )
1000
1001    # i-direction: f-point segment direction (around v-point)
1002    zlam = glamf 
1003    zphi = gphif
1004    zlan = np.roll (glamf, axis=-1, shift=1) # glamf (ji-1,jj)
1005    zphh = np.roll (gphif, axis=-1, shift=1) # gphif (ji-1,jj)
1006    zxffv =  2.0 * np.cos ( rad*zlam ) * np.tan ( rpi/4. - rad*zphi/2. )   \
1007          -  2.0 * np.cos ( rad*zlan ) * np.tan ( rpi/4. - rad*zphh/2. )
1008    zyffv =  2.0 * np.sin ( rad*zlam ) * np.tan ( rpi/4. - rad*zphi/2. )   \
1009          -  2.0 * np.sin ( rad*zlan ) * np.tan ( rpi/4. - rad*zphh/2. )
1010    znffv = np.sqrt ( znnpv * ( zxffv*zxffv + zyffv*zyffv )  )
1011
1012    # j-direction: u-point segment direction (around f-point)
1013    zlam = np.roll (glamu, axis=-2, shift=-1) # glamu (ji,jj+1)
1014    zphi = np.roll (gphiu, axis=-2, shift=-1) # gphiu (ji,jj+1)
1015    zlan = glamu
1016    zphh = gphiu
1017    zxuuf =  2. * np.cos ( rad*zlam ) * np.tan ( rpi/4. - rad*zphi/2. )   \
1018          -  2. * np.cos ( rad*zlan ) * np.tan ( rpi/4. - rad*zphh/2. )
1019    zyuuf =  2. * np.sin ( rad*zlam ) * np.tan ( rpi/4. - rad*zphi/2. )   \
1020          -  2. * np.sin ( rad*zlan ) * np.tan ( rpi/4. - rad*zphh/2. )
1021    znuuf = np.sqrt ( znnpf * ( zxuuf*zxuuf + zyuuf*zyuuf )  )
1022
1023   
1024    # cosinus and sinus using scalar and vectorial products
1025    gsint = ( zxnpt*zyvvt - zynpt*zxvvt ) / znvvt
1026    gcost = ( zxnpt*zxvvt + zynpt*zyvvt ) / znvvt
1027   
1028    gsinu = ( zxnpu*zyffu - zynpu*zxffu ) / znffu
1029    gcosu = ( zxnpu*zxffu + zynpu*zyffu ) / znffu
1030   
1031    gsinf = ( zxnpf*zyuuf - zynpf*zxuuf ) / znuuf
1032    gcosf = ( zxnpf*zxuuf + zynpf*zyuuf ) / znuuf
1033   
1034    gsinv = ( zxnpv*zxffv + zynpv*zyffv ) / znffv
1035    gcosv =-( zxnpv*zyffv - zynpv*zxffv ) / znffv  # (caution, rotation of 90 degres)
1036   
1037    gsint = lbc (gsint, cd_type='T', nperio=nperio, psgn=-1.)
1038    gcost = lbc (gcost, cd_type='T', nperio=nperio, psgn=-1.)
1039    gsinu = lbc (gsinu, cd_type='U', nperio=nperio, psgn=-1.)
1040    gcosu = lbc (gcosu, cd_type='U', nperio=nperio, psgn=-1.)
1041    gsinv = lbc (gsinv, cd_type='V', nperio=nperio, psgn=-1.)
1042    gcosv = lbc (gcosv, cd_type='V', nperio=nperio, psgn=-1.)
1043    gsinf = lbc (gsinf, cd_type='F', nperio=nperio, psgn=-1.)
1044    gcosf = lbc (gcosf, cd_type='F', nperio=nperio, psgn=-1.)
1045
1046    if math == xr :
1047        gsint = gsint.assign_coords ( glamt.coords )
1048        gcost = gcost.assign_coords ( glamt.coords )
1049        gsinu = gsinu.assign_coords ( glamu.coords )
1050        gcosu = gcosu.assign_coords ( glamu.coords )
1051        gsinv = gsinv.assign_coords ( glamv.coords )
1052        gcosv = gcosv.assign_coords ( glamv.coords )
1053        gsinf = gsinf.assign_coords ( glamf.coords )
1054        gcosf = gcosf.assign_coords ( glamf.coords )
1055
1056    return gsint, gcost, gsinu, gcosu, gsinv, gcosv, gsinf, gcosf
1057
1058def angle (glam, gphi, nperio, cd_type='T') :
1059    '''Compute sinus and cosinus of model line direction with respect to east'''
1060    math = __math__ (glam) 
1061    # north pole direction & modulous
1062    zxnp = 0. - 2.0 * np.cos (rad*glam) * np.tan (rpi/4.0 - rad*gphi/2.0)
1063    zynp = 0. - 2.0 * np.sin (rad*glam) * np.tan (rpi/4.0 - rad*gphi/2.0)
1064    znnp = zxnp*zxnp + zynp*zynp
1065
1066    # j-direction: segment direction (around point)
1067    zlan_n = np.roll (glam, axis=-2, shift=-1) # glam [jj+1, ji]
1068    zphh_n = np.roll (gphi, axis=-2, shift=-1) # gphi [jj+1, ji]
1069    zlan_s = np.roll (glam, axis=-2, shift= 1) # glam [jj-1, ji]
1070    zphh_s = np.roll (gphi, axis=-2, shift= 1) # gphi [jj-1, ji]
1071   
1072    zxff = 2.0 * np.cos (rad*zlan_n) * np.tan (rpi/4.0 - rad*zphh_n/2.0) \
1073        -  2.0 * np.cos (rad*zlan_s) * np.tan (rpi/4.0 - rad*zphh_s/2.0)
1074    zyff = 2.0 * np.sin (rad*zlan_n) * np.tan (rpi/4.0 - rad*zphh_n/2.0) \
1075        -  2.0 * np.sin (rad*zlan_s) * np.tan (rpi/4.0 - rad*zphh_s/2.0)
1076    znff = np.sqrt (znnp * (zxff*zxff + zyff*zyff) )
1077 
1078    gsin = (zxnp*zyff - zynp*zxff) / znff
1079    gcos = (zxnp*zxff + zynp*zyff) / znff
1080
1081    gsin = lbc (gsin, cd_type=cd_type, nperio=nperio, psgn=-1.)
1082    gcos = lbc (gcos, cd_type=cd_type, nperio=nperio, psgn=-1.)
1083
1084    if math == xr :
1085        gsin = gsin.assign_coords ( glam.coords )
1086        gcos = gcos.assign_coords ( glam.coords )
1087       
1088    return gsin, gcos
1089
1090def rot_en2ij ( u_e, v_n, gsin, gcos, nperio, cd_type ) :
1091    '''
1092    ** Purpose :   Rotate the Repere: Change vector componantes between
1093    geographic grid --> stretched coordinates grid.
1094    All components are on the same grid (T, U, V or F)
1095    '''
1096
1097    u_i = + u_e * gcos + v_n * gsin
1098    v_j = - u_e * gsin + v_n * gcos
1099   
1100    u_i = lbc (u_i, nperio=nperio, cd_type=cd_type, psgn=-1.0)
1101    v_j = lbc (v_j, nperio=nperio, cd_type=cd_type, psgn=-1.0)
1102   
1103    return u_i, v_j
1104
1105def rot_ij2en ( u_i, v_j, gsin, gcos, nperio, cd_type='T' ) :
1106    '''
1107    ** Purpose :   Rotate the Repere: Change vector componantes from
1108    stretched coordinates grid --> geographic grid
1109    All components are on the same grid (T, U, V or F)
1110    '''
1111    u_e = + u_i * gcos - v_j * gsin
1112    v_n = + u_i * gsin + v_j * gcos
1113   
1114    u_e = lbc (u_e, nperio=nperio, cd_type=cd_type, psgn= 1.0)
1115    v_n = lbc (v_n, nperio=nperio, cd_type=cd_type, psgn= 1.0)
1116   
1117    return u_e, v_n
1118
1119def rot_uv2en ( uo, vo, gsint, gcost, nperio, zdim='deptht' ) :
1120    '''
1121    ** Purpose :   Rotate the Repere: Change vector componantes from
1122    stretched coordinates grid --> geographic grid
1123    uo is on the U grid point, vo is on the V grid point
1124    east-north components on the T grid point   
1125    '''
1126    math = __math__ (uo)
1127
1128    ut = U2T (uo, nperio=nperio, psgn=-1.0, zdim=zdim)
1129    vt = V2T (vo, nperio=nperio, psgn=-1.0, zdim=zdim)
1130   
1131    u_e = + ut * gcost - vt * gsint
1132    v_n = + ut * gsint + vt * gcost
1133
1134    u_e = lbc (u_e, nperio=nperio, cd_type='T', psgn=1.0)
1135    v_n = lbc (v_n, nperio=nperio, cd_type='T', psgn=1.0)
1136   
1137    return u_e, v_n
1138
1139def rot_uv2enF ( uo, vo, gsinf, gcosf, nperio, zdim='deptht' ) :
1140    '''
1141    ** Purpose :   Rotate the Repere: Change vector componantes from
1142    stretched coordinates grid --> geographic grid
1143    uo is on the U grid point, vo is on the V grid point
1144    east-north components on the T grid point   
1145    '''
1146    math = __math__ (uo)
1147
1148    uf = U2F (uo, nperio=nperio, psgn=-1.0, zdim=zdim)
1149    vf = V2F (vo, nperio=nperio, psgn=-1.0, zdim=zdim)
1150   
1151    u_e = + uf * gcosf - vf * gsinf
1152    v_n = + uf * gsinf + vf * gcosf
1153
1154    u_e = lbc (u_e, nperio=nperio, cd_type='F', psgn= 1.0)
1155    v_n = lbc (v_n, nperio=nperio, cd_type='F', psgn= 1.0)
1156   
1157    return u_e, v_n
1158             
1159def U2T (utab, nperio=None, psgn=-1.0, zdim='deptht') :
1160    '''Interpolate an array from U grid to T grid i-mean)'''
1161    math = __math__ (utab)
1162    utab_0 = math.where ( np.isnan(utab), 0., utab)
1163    utab_0 = lbc (utab_0 , nperio=nperio, cd_type='U', psgn=psgn)
1164    ix, ax = __findAxis__ (utab_0, 'x')
1165    iz, az = __findAxis__ (utab_0, 'z')
1166    ttab =  lbc ( 0.5 * (utab_0 + np.roll (utab_0, axis=ix, shift=1)), cd_type='T', nperio=nperio, psgn=psgn)
1167   
1168    if math == xr :
1169        if ax != None :
1170            ttab = ttab.assign_coords({ax:np.arange (ttab.shape[ix])+1.})
1171        if zdim != None and iz != None  and az != 'olevel' : 
1172            ttab = ttab.rename( {az:zdim}) 
1173    return ttab
1174
1175def V2T (vtab, nperio=None, psgn=-1.0, zdim='deptht') :
1176    '''Interpolate an array from V grid to T grid (j-mean)'''
1177    math = __math__ (vtab)
1178    vtab_0 = math.where ( np.isnan(vtab), 0., vtab)
1179    vtab_0 = lbc (vtab_0 , nperio=nperio, cd_type='V', psgn=psgn)
1180    iy, ay = __findAxis__ (vtab_0, 'y')
1181    iz, az = __findAxis__ (vtab_0, 'z')
1182    ttab = lbc ( 0.5 * (vtab_0 + np.roll (vtab_0, axis=iy, shift=1)), cd_type='T', nperio=nperio, psgn=psgn)
1183                     
1184    if math == xr :
1185        if ay !=None : 
1186            ttab = ttab.assign_coords({ay:np.arange(ttab.shape[iy])+1.})
1187        if zdim != None and iz != None  and az != 'olevel' :
1188            ttab = ttab.rename( {az:zdim}) 
1189    return ttab
1190
1191def F2T (ftab, nperio=None, psgn=1.0, zdim='depthf') :
1192    '''Interpolate an array from F grid to T grid (i- and j- means)'''
1193    math = __math__ (ftab)
1194    ftab_0 = math.where ( np.isnan(ftab), 0., ftab)
1195    ftab_0 = lbc (ftab_0 , nperio=nperio, cd_type='F', psgn=psgn)
1196    ttab = V2T(F2V(ftab_0, nperio=nperio, psgn=psgn, zdim=zdim), nperio=nperio, psgn=psgn, zdim=zdim)
1197    return ttab
1198
1199def T2U (ttab, nperio=None, psgn=1.0, zdim='depthu') :
1200    '''Interpolate an array from T grid to U grid (i-mean)'''
1201    math = __math__ (ttab)
1202    ttab_0 = math.where ( np.isnan(ttab), 0., ttab)
1203    ttab_0 = lbc (ttab_0 , nperio=nperio, cd_type='T', psgn=psgn)
1204    ix, ax = __findAxis__ (ttab_0, 'x')
1205    iz, az = __findAxis__ (ttab_0, 'z')
1206    utab = lbc ( 0.5 * (ttab_0 + np.roll (ttab_0, axis=ix, shift=-1)), cd_type='U', nperio=nperio, psgn=psgn)
1207
1208    if math == xr :   
1209        if ax != None : 
1210            utab = ttab.assign_coords({ax:np.arange(utab.shape[ix])+1.})
1211        if zdim != None  and iz != None  and az != 'olevel' :
1212            utab = utab.rename( {az:zdim}) 
1213    return utab
1214
1215def T2V (ttab, nperio=None, psgn=1.0, zdim='depthv') :
1216    '''Interpolate an array from T grid to V grid (j-mean)'''
1217    math = __math__ (ttab)
1218    ttab_0 = math.where ( np.isnan(ttab), 0., ttab)
1219    ttab_0 = lbc (ttab_0 , nperio=nperio, cd_type='T', psgn=psgn)
1220    iy, ay = __findAxis__ (ttab_0, 'y')
1221    iz, az = __findAxis__ (ttab_0, 'z')
1222    vtab = lbc ( 0.5 * (ttab_0 + np.roll (ttab_0, axis=iy, shift=-1)), cd_type='V', nperio=nperio, psgn=psgn)
1223
1224    if math == xr :
1225        if ay != None : 
1226            vtab = vtab.assign_coords({ay:np.arange(vtab.shape[iy])+1.})
1227        if zdim != None  and iz != None and az != 'olevel' :
1228            vtab = vtab.rename( {az:zdim}) 
1229    return vtab
1230
1231def V2F (vtab, nperio=None, psgn=-1.0, zdim='depthf') :
1232    '''Interpolate an array from V grid to F grid (i-mean)'''
1233    math = __math__ (vtab)
1234    vtab_0 = math.where ( np.isnan(vtab), 0., vtab)
1235    vtab_0 = lbc (vtab_0 , nperio=nperio, cd_type='V', psgn=psgn)
1236    ix, ax = __findAxis__ (vtab_0, 'x')
1237    iz, az = __findAxis__ (vtab_0, 'z')
1238    ftab = lbc ( 0.5 * (vtab_0 + np.roll (vtab_0, axis=ix, shift=-1)), cd_type='F', nperio=nperio, psgn=psgn)
1239
1240    if math == xr :
1241        if ax != None : 
1242            ftab = ftab.assign_coords({ax:np.arange(ftab.shape[ix])+1.})
1243        if zdim != None and iz != None and az != 'olevel' :
1244            ftab = ftab.rename( {az:zdim}) 
1245    return ftab
1246
1247def U2F (utab, nperio=None, psgn=-1.0, zdim='depthf') :
1248    '''Interpolate an array from U grid to F grid i-mean)'''
1249    math = __math__ (utab)
1250    utab_0 = math.where ( np.isnan(utab), 0., utab)
1251    utab_0 = lbc (utab_0 , nperio=nperio, cd_type='U', psgn=psgn)
1252    iy, ay = __findAxis__ (utab_0, 'y')
1253    iz, az = __findAxis__ (utab_0, 'z')
1254    ftab = lbc ( 0.5 * (utab_0 + np.roll (utab_0, axis=iy, shift=-1)), cd_type='F', nperio=nperio, psgn=psgn)
1255
1256    if math == xr :
1257        if ay != None : 
1258            ftab = ftab.assign_coords({'y':np.arange(ftab.shape[iy])+1.})
1259        if zdim != None and iz != None and az != 'olevel' :
1260            ftab = ftab.rename( {az:zdim}) 
1261    return ftab
1262
1263def F2T (ftab, nperio=None, psgn=1.0, zdim='deptht') :
1264    '''Interpolate an array on F grid to T grid (i- and j- means)'''
1265    math = __math__ (ftab)
1266    ftab_0 = math.where ( np.isnan(ttab), 0., ttab)
1267    ftab_0 = lbc (ftab_0 , nperio=nperio, cd_type='F', psgn=psgn)
1268    ttab = U2T(F2U(ftab_0, nperio=nperio, psgn=psgn, zdim=zdim), nperio=nperio, psgn=psgn, zdim=zdim)
1269    return ttab
1270
1271def T2F (ttab, nperio=None, psgn=1.0, zdim='deptht') :
1272    '''Interpolate an array on T grid to F grid (i- and j- means)'''
1273    math = __math__ (ftab)
1274    ttab_0 = math.where ( np.isnan(ttab), 0., ttab)
1275    ttab_0 = lbc (ttab_0 , nperio=nperio, cd_type='T', psgn=psgn)
1276    ftab = T2U(U2F(ttab, nperio=nperio, psgn=psgn, zdim=zdim), nperio=nperio, psgn=psgn, zdim=zdim)
1277    return ftab
1278
1279def F2U (ftab, nperio=None, psgn=1.0, zdim='depthu') :
1280    '''Interpolate an array on F grid to FUgrid (i-mean)'''
1281    math = __math__ (ftab)
1282    ftab_0 = math.where ( np.isnan(ftab), 0., ftab)
1283    ftab_0 = lbc (ftab_0 , nperio=nperio, cd_type='F', psgn=psgn)
1284    iy, ay = __findAxis__ (ftab_0, 'y')
1285    iz, az = __findAxis__ (ftab_0, 'z')
1286    utab = lbc ( 0.5 * (ftab_0 + np.roll (ftab_0, axis=iy, shift=-1)), cd_type='U', nperio=nperio, psgn=psgn)
1287   
1288    if math == xr :
1289        utab = utab.assign_coords({ay:np.arange(ftab.shape[iy])+1.})
1290        if zdim != None and iz != None and az != 'olevel' :
1291            utab = utab.rename( {az:zdim}) 
1292    return utab
1293
1294def F2V (ftab, nperio=None, psgn=1.0, zdim='depthv') :
1295    '''Interpolate an array from F grid to V grid (i-mean)'''
1296    math = __math__ (ftab)
1297    ftab_0 = math.where ( np.isnan(ftab), 0., ftab)
1298    ftab_0 = lbc (ftab_0 , nperio=nperio, cd_type='F', psgn=psgn)
1299    ix, ax = __findAxis__ (ftab_0, 'x')
1300    iz, az = __findAxis__ (ftab_0, 'z')
1301    vtab = lbc ( 0.5 * (ftab_0 + np.roll (ftab_0, axis=ix, shift=-1)), cd_type='V', nperio=nperio)
1302
1303    if math == xr :
1304        vtab = vtab.assign_coords({ax:np.arange(ftab.shape[ix])+1.})
1305        if zdim != None and iz != None and az != 'olevel' :
1306            vtab = vtab.rename( {az:zdim}) 
1307    return vtab
1308
1309def W2T (wtab, zcoord=None, zdim='deptht', sval=np.nan) :
1310    '''
1311    Interpolate an array on W grid to T grid (k-mean)
1312    sval is the bottom value
1313    '''
1314    math = __math__ (wtab)
1315    wtab_0 = math.where ( np.isnan(wtab), 0., wtab)
1316
1317    iz, az = __findAxis__ (wtab_0, 'z')
1318       
1319    ttab = 0.5 * ( wtab_0 + np.roll (wtab_0, axis=iz, shift=-1) )
1320   
1321    if math == xr :
1322        ttab[{az:iz}] = sval
1323        if zdim != None and iz != None and az != 'olevel' :
1324            ttab = ttab.rename ( {az:zdim} )
1325        try    : ttab = ttab.assign_coords ( {zdim:zcoord} )
1326        except : pass
1327    else :
1328        ttab[..., -1, :, :] = sval
1329
1330    return ttab
1331
1332def T2W (ttab, zcoord=None, zdim='depthw', sval=np.nan, extrap_surf=False) :
1333    '''Interpolate an array from T grid to W grid (k-mean)
1334    sval is the surface value
1335    if extrap_surf==True, surface value is taken from 1st level value.
1336    '''
1337    math = __math__ (ttab)
1338    ttab_0 = math.where ( np.isnan(ttab), 0., ttab)
1339    iz, az = __findAxis__ (ttab_0, 'z')
1340    wtab = 0.5 * ( ttab_0 + np.roll (ttab_0, axis=iz, shift=1) )
1341
1342    if math == xr :
1343        if extrap_surf : wtab[{az:0}] = ttabb[{az:0}]
1344        else           : wtab[{az:0}] = sval
1345    else : 
1346        if extrap_surf : wtab[..., 0, :, :] = ttab[..., 0, :, :]
1347        else           : wtab[..., 0, :, :] = sval
1348
1349    if math == xr :
1350        if zdim != None and iz != None and az != 'olevel' :
1351                wtab = wtab.rename ( {az:zdim})
1352        if zcoord != None : wtab = wtab.assign_coords ( {zdim:zcoord})
1353        else              : ztab = wtab.assign_coords ( {zdim:np.arange(ttab.shape[iz])+1.} )
1354    return wtab
1355
1356def fill (ptab, nperio, cd_type='T', npass=1) :
1357    '''
1358    Fill 0. or np.nan values with mean of neighbours
1359   
1360    Inputs :
1361       ptab : input field to fill
1362       nperio, cd_type : periodicity characteristics
1363    '''       
1364
1365    math = __math__ (ptab)
1366    ztab   = math.where (np.isnan(ptab), 0., ptab)
1367
1368    DoPerio = False
1369    if nperio == 4.2 or nperio == 6.2 : DoPerio = True
1370       
1371    if DoPerio :  ztab = lbc_add (ztab)
1372   
1373    for nn in np.arange (npass) : 
1374        zmask  = math.where (ztab==0., 0., 1.  )
1375        # Compte du nombre de voisins
1376        zcount = zmask \
1377          + np.roll(zmask, shift=1, axis=-1) + np.roll(zmask, shift=-1, axis=-1) \
1378          + np.roll(zmask, shift=1, axis=-2) + np.roll(zmask, shift=-1, axis=-2) \
1379          + 0.5 * ( \
1380                + np.roll(np.roll(zmask, shift= 1, axis=-2), shift= 1, axis=-1) \
1381                + np.roll(np.roll(zmask, shift=-1, axis=-2), shift= 1, axis=-1) \
1382                + np.roll(np.roll(zmask, shift= 1, axis=-2), shift=-1, axis=-1) \
1383                + np.roll(np.roll(zmask, shift=-1, axis=-2), shift=-1, axis=-1) )
1384
1385        znew = zmask \
1386          + np.roll(ztab, shift=1, axis=-1) + np.roll(ztab, shift=-1, axis=-1) \
1387          + np.roll(ztab, shift=1, axis=-2) + np.roll(ztab, shift=-1, axis=-2) \
1388          + 0.5 * ( \
1389            + np.roll(np.roll(ztab, shift= 1, axis=-2), shift= 1, axis=-1) \
1390            + np.roll(np.roll(ztab, shift=-1, axis=-2), shift= 1, axis=-1) \
1391            + np.roll(np.roll(ztab, shift= 1, axis=-2), shift=-1, axis=-1) \
1392            + np.roll(np.roll(ztab, shift=-1, axis=-2), shift=-1, axis=-1) )
1393
1394        zcount = lbc (zcount, nperio=nperio, cd_type=cd_type)
1395        znew   = lbc (znew  , nperio=nperio, cd_type=cd_type)
1396   
1397        ztab = math.where (np.logical_and (zmask==0., zcount>0), znew/zcount, ztab)
1398       
1399    ztab = math.where (ztab==0., np.nan, ztab)
1400
1401    if DoPerio : ztab = lbc_del (ztab)
1402
1403    return ztab
1404
1405def correct_uv (u, v, lat) :
1406    '''
1407    Correct a Cartopy bug in Orthographic projection
1408
1409    See https://github.com/SciTools/cartopy/issues/1179
1410
1411    The correction is needed with cartopy <= 0.20
1412    It seems that version 0.21 will correct the bug (https://github.com/SciTools/cartopy/pull/1926)
1413
1414    Inputs :
1415       u, v : eastward/nothward components
1416       lat  : latitude of the point (degrees north)
1417
1418    Outputs :
1419       modified eastward/nothward components have correct polar projection in cartopy
1420    '''
1421    math = __math__ (u)
1422    uv = np.sqrt (u*u + v*v)           # Original modulus
1423    zu = u
1424    zv = v * np.cos (rad*lat)
1425    zz = np.sqrt ( zu*zu + zv*zv )     # Corrected modulus
1426    uc = zu*uv/zz ; vc = zv*uv/zz      # Final corrected values
1427    return uc, vc
1428
1429## ===========================================================================
1430##
1431##                               That's all folk's !!!
1432##
1433## ===========================================================================
1434
1435def is_orca_north_fold ( Xtest, cname_long='T' ) :
1436    '''
1437    Ported (pirated !!?) from Sosie
1438
1439    Tell if there is a 2/point band overlaping folding at the north pole typical of the ORCA grid
1440
1441    0 => not an orca grid (or unknown one)
1442    4 => North fold T-point pivot (ex: ORCA2)
1443    6 => North fold F-point pivot (ex: ORCA1)
1444
1445    We need all this 'cname_long' stuff because with our method, there is a
1446    confusion between "Grid_U with T-fold" and "Grid_V with F-fold"
1447    => so knowing the name of the longitude array (as in namelist, and hence as
1448    in netcdf file) might help taking the righ decision !!! UGLY!!!
1449    => not implemented yet
1450    '''
1451   
1452    ifld_nord =  0 ; cgrd_type = 'X'
1453    ny, nx = Xtest.shape[-2:]
1454
1455    if ny > 3 : # (case if called with a 1D array, ignoring...)
1456        if ( Xtest [ny-1, 1:nx//2-1] - Xtest [ny-3, nx-1:nx-nx//2+1:-1] ).sum() == 0. :
1457          ifld_nord = 4 ; cgrd_type = 'T' # T-pivot, grid_T     
1458
1459        if ( Xtest [ny-1, 1:nx//2-1] - Xtest [ny-3, nx-2:nx-nx//2  :-1] ).sum() == 0. :
1460            if cnlon == 'U' : ifld_nord = 4 ;  cgrd_type = 'U' # T-pivot, grid_T
1461                ## LOLO: PROBLEM == 6, V !!!
1462
1463        if ( Xtest [ny-1, 1:nx//2-1] - Xtest [ny-3, nx-1:nx-nx//2+1:-1] ).sum() == 0. :
1464            ifld_nord = 4 ; cgrd_type = 'V' # T-pivot, grid_V
1465
1466        if ( Xtest [ny-1, 1:nx//2-1] - Xtest [ny-2, nx-1-1:nx-nx//2:-1] ).sum() == 0. :
1467            ifld_nord = 6 ; cgrd_type = 'T'# F-pivot, grid_T
1468
1469        if ( Xtest [ny-1, 1:nx//2-1] - Xtest [ny-1, nx-1:nx-nx//2-1:-1] ).sum() == 0. :
1470            ifld_nord = 6 ;  cgrd_type = 'U' # F-pivot, grid_U
1471
1472        if ( Xtest [ny-1, 1:nx//2-1] - Xtest [ny-3, nx-2:nx-nx//2  :-1] ).sum() == 0. :
1473            if cnlon == 'V' : ifld_nord = 6 ; cgrd_type = 'V' # F-pivot, grid_V
1474                ## LOLO: PROBLEM == 4, U !!!
1475
1476    return ifld_nord, cgrd_type
Note: See TracBrowser for help on using the repository browser.