# -*- coding: utf-8 -*- ## =========================================================================== ## ## Warning, to install, configure, run, use any of Olivier Marti's ## software or to read the associated documentation you'll need at least ## one (1) brain in a reasonably working order. Lack of this implement ## will void any warranties (either express or implied). ## O. Marti assumes no responsability for errors, omissions, ## data loss, or any other consequences caused directly or indirectly by ## the usage of his software by incorrectly or partially configured ## personal. ## ''' Utilities to plot NEMO ORCA fields Periodicity and other stuff olivier.marti@lsce.ipsl.fr ''' import sys, numpy as np try : import xarray as xr except : pass rpi = np.pi ; rad = rpi / 180.0 nperio_valid_range = [0,1,4,5,6] ## SVN information __Author__ = "$Author$" __Date__ = "$Date$" __Revision__ = "$Revision$" __Id__ = "$Id$" __HeadURL = "$HeadURL$" def __guessNperio__ (jpj, jpi, nperio=None) : ''' Tries to guess the value of nperio (periodicity parameter. See NEMO documentation for details) Inputs jpi : number of longitudes nperio : periodicity parameter ''' if nperio == None : if jpi == 182 : nperio = 4 # ORCA2. We choose legacy orca2. if jpi == 362 : nperio = 6 # eORCA1. if jpi == 1442 : nperio = 6 # ORCA025. if jpj == 149 : nperio = 4 # ORCA2. We choose legacy orca2. if jpj == 332 : nperio = 6 # eORCA1. # if nperio == None : sys.exit ('in nemo module : nperio not found, and cannot by guessed' ) else : print ('nperio set as {:d} (deduced from jpi={:d})'.format (nperio, jpi)) return nperio def __guessPoint__ (ptab) : ''' Tries to guess the grid point (periodicity parameter. See NEMO documentation for details) For array conforments with xgcm requirements Inputs ptab : xarray array Credits : who is the original author ? ''' gP = None if isinstance (ptab, xr.core.dataarray.DataArray) : if 'x_c' in ptab.dims and 'y_c' in ptab.dims : gP = 'T' if 'x_f' in ptab.dims and 'y_c' in ptab.dims : gP = 'U' if 'x_c' in ptab.dims and 'y_f' in ptab.dims : gP = 'V' if 'x_f' in ptab.dims and 'y_f' in ptab.dims : gP = 'F' if 'x_c' in ptab.dims and 'y_c' in ptab.dims and 'z_c' in ptab.dims : gP = 'T' if 'x_c' in ptab.dims and 'y_c' in ptab.dims and 'z_f' in ptab.dims : gP = 'W' if 'x_f' in ptab.dims and 'y_c' in ptab.dims and 'z_f' in ptab.dims : gP = 'U' if 'x_c' in ptab.dims and 'y_f' in ptab.dims and 'z_f' in ptab.dims : gP = 'V' if 'x_f' in ptab.dims and 'y_f' in ptab.dims and 'z_f' in ptab.dims : gP = 'F' if gP == None : sys.exit ('in nemo module : cd_type not found, and cannot by guessed' ) else : print ('Grid set as', gP, 'deduced from dims ', ptab.dims) return gP else : sys.exit ('in nemo module : cd_type not found, input is not an xarray data' ) #return gP def fixed_lon (lon) : ''' Returns corrected longitudes for nicer plots fixed_lon (lon) lon : longitudes of the grid. At least 2D. Designed by Phil Pelson. See https://gist.github.com/pelson/79cf31ef324774c97ae7 ''' fixed_lon = lon.copy () for i, start in enumerate (np.argmax (np.abs (np.diff (lon, axis=-1)) > 180, axis=-1)) : fixed_lon[..., i, start+1:] += 360 # Special case for eORCA025 if lon.shape[-1] == 1442 : lon [..., -2, :] = lon [..., -3, :] return fixed_lon def jeq (lat) : ''' Returns j index of equator in the grid lat : latitudes of the grid. At least 2D. ''' jeq = np.argmin (np.abs (np.float64(lat[:,0]))) return jeq def lon1D (lon, lat=None) : ''' Returns 1D longitude for simple plots. lon : longitudes of the grid lat (optionnal) : latitudes of the grid ''' if np.max (lat) != None : je = jeq (lat) lon1D = lon[..., je, :] else : jpj, jpi = lon.shape[-2:] lon1D = lon[..., jpj//3, :] start = np.argmax (np.abs (np.diff (lon1D, axis=-1)) > 180.0, axis=-1) lon1D[..., start+1:] += 360 return lon1D def latreg (lat, diff=0.1) : ''' Returns maximum j index where gridlines are along latitude in the northern hemisphere lat : latitudes of the grid (2D) diff [optional] : tolerance ''' if diff == None : dy = np.float64 (np.mean (np.abs (lat - np.roll(lat,shift=1,axis=-2)))) diff = dy/100. je = jeq (lat) jreg = np.where (lat[...,je:,:].max(axis=-1) - lat[...,je:,:].min(axis=-1)< diff)[-1][-1] + je latreg = np.float64 (lat[...,jreg,:].mean(axis=-1)) JREG = jreg return jreg, latreg def lat1D (lat) : ''' Returns 1D latitudes for zonal means and simple plots. lat : latitudes of the grid (2D) ''' jpj, jpi = lat.shape[-2:] try : if type (lat) == xr.core.dataarray.DataArray : math = xr else : math = np except : math = np dy = np.float64 (np.mean (np.abs (lat - np.roll (lat, shift=1,axis=-2)))) je = jeq (lat) lat_eq = np.float64 (lat[...,je,:].mean(axis=-1)) jreg, lat_reg = latreg (lat) lat_ave = np.mean (lat, axis=-1) if ( np.abs (lat_eq) < dy/100. ) : # T, U or W grid dys = (90.-lat_reg) / (jpj-jreg-1)*0.5 yrange = 90.-dys-lat_reg else : # V or F grid yrange = 90. -lat_reg lat1D = math.where (lat_avey0, latx0, lonx0, lon+360x0, lon-360 jpi : extend = tab else : if nperio == 1 : istart = 0 ; le=jpi+1 ; la=0 if nperio > 1 : # OPA case with to halo points for periodicity istart = 1 ; le=jpi-2 ; la=1 # Perfect, except at the pole that should be masked by lbc_plot if math == xr : extend = xr.concat ( (tab[..., istart:istart+le+1] + xplus, tab[..., istart+la:jplus]), dim=tab.dims[-1] ) if math == np : extend = np.concatenate ( (tab[..., istart:istart+le+1] + xplus, tab[..., istart+la:jplus]), axis=-1 ) return extend def orca2reg (ff, lat_name='nav_lat', lon_name='nav_lon', y_name='y', x_name='x') : ''' Assign ORCA dataset on a regular grid. For use in the tropical region. Inputs : ff : xarray dataset lat_name, lon_name : name of latitude and longitude 2D field in ff y_name, x_name : namex of dimensions in ff Returns : xarray dataset with rectangular grid. Incorrect above 20°N ''' # Compute 1D longitude and latitude (lat, lon) = latlon1D (ff[lat_name], ff[lon_name]) # Assign lon and lat as dimensions of the dataset if y_name in ff.dims : lat = xr.DataArray (lat, coords=[lat,], dims=['lat',]) ff = ff.rename_dims ({y_name: "lat",}).assign_coords (lat=lat) if x_name in ff.dims : lon = xr.DataArray (lon, coords=[lon,], dims=['lon',]) ff = ff.rename_dims ({x_name: "lon",}).assign_coords (lon=lon) # Force dimensions to be in the right order coord_order = ['lat', 'lon'] for dim in [ 'depthw', 'depthv', 'depthu', 'deptht', 'depth', 'z', 'time_counter', 'time', 'tbnds', 'bnds', 'axis_nbounds', 'two2', 'two1', 'two', 'four',] : if dim in ff.dims : coord_order.insert (0, dim ) ff = ff.transpose (*coord_order) return ff def lbc_init (ptab, nperio=None, cd_type='T') : """ Prepare for all lbc calls Set periodicity on input field nperio : Type of periodicity 0 : No periodicity 1, 4, 6 : Cyclic on i dimension (generaly longitudes) with 2 points halo 2 : Obsolete (was symmetric condition at southern boundary ?) 3, 4 : North fold T-point pivot (legacy ORCA2) 5, 6 : North fold F-point pivot (ORCA1, ORCA025, ORCA2 with new grid for paleo) cd_type : Grid specification : T, U, V or F See NEMO documentation for further details """ jpj, jpi = ptab.shape[-2:] if nperio == None : nperio = __guessNperio__ (jpj, jpi, nperio) if nperio not in nperio_valid_range : print ( 'nperio=', nperio, ' is not in the valid range', nperio_valid_range ) sys.exit() if cd_type == None and isinstance (ptab, xr.core.dataarray.DataArray) : cd_type = __guessPoint__ (ptab) if cd_type not in ['T', 'U', 'V', 'F', 'W', 'F' ]: print ( 'cd_type=', cd_type, ' is not in the list T, U, V, F, W, F' ) sys.exit () return jpj, jpi, nperio, cd_type def lbc (ptab, nperio=None, cd_type='T', psgn=1.0, nemo_4U_bug=False) : """ Set periodicity on input field ptab : Input array (works rank 2 at least : ptab[...., lat, lon] nperio : Type of periodicity cd_type : Grid specification : T, U, V or F psgn : For change of sign for vector components (1 for scalars, -1 for vector components) See NEMO documentation for further details """ jpj, jpi, nperio, cd_type = lbc_init (ptab, nperio, cd_type) psgn = ptab.dtype.type (psgn) # #> East-West boundary conditions # ------------------------------ if nperio in [1, 4, 6] : # ... cyclic ptab [..., :, 0] = ptab [..., :, -2] ptab [..., :, -1] = ptab [..., :, 1] # #> North-South boundary conditions # -------------------------------- if nperio in [3, 4] : # North fold T-point pivot if cd_type in [ 'T', 'W' ] : # T-, W-point ptab [..., -1, 1: ] = psgn * ptab [..., -3, -1:0:-1] ptab [..., -1, 0 ] = psgn * ptab [..., -3, 2 ] ptab [..., -2, jpi//2: ] = psgn * ptab [..., -2, jpi//2:0:-1 ] if cd_type == 'U' : ptab [..., -1, 0:-1 ] = psgn * ptab [..., -3, -1:0:-1 ] ptab [..., -1, 0 ] = psgn * ptab [..., -3, 1 ] ptab [..., -1, -1 ] = psgn * ptab [..., -3, -2 ] if nemo_4U_bug : ptab [..., -2, jpi//2+1:-1] = psgn * ptab [..., -2, jpi//2-2:0:-1] ptab [..., -2, jpi//2-1 ] = psgn * ptab [..., -2, jpi//2 ] else : ptab [..., -2, jpi//2-1:-1] = psgn * ptab [..., -2, jpi//2:0:-1] if cd_type == 'V' : ptab [..., -2, 1: ] = psgn * ptab [..., -3, jpi-1:0:-1 ] ptab [..., -1, 1: ] = psgn * ptab [..., -4, -1:0:-1 ] ptab [..., -1, 0 ] = psgn * ptab [..., -4, 2 ] if cd_type == 'F' : ptab [..., -2, 0:-1 ] = psgn * ptab [..., -3, -1:0:-1 ] ptab [..., -1, 0:-1 ] = psgn * ptab [..., -4, -1:0:-1 ] ptab [..., -1, 0 ] = psgn * ptab [..., -4, 1 ] ptab [..., -1, -1 ] = psgn * ptab [..., -4, -2 ] if nperio in [5, 6] : # North fold F-point pivot if cd_type in ['T', 'W'] : ptab [..., -1, 0: ] = psgn * ptab [..., -2, -1::-1 ] if cd_type == 'U' : ptab [..., -1, 0:-1 ] = psgn * ptab [..., -2, -2::-1 ] ptab [..., -1, -1 ] = psgn * ptab [..., -2, 0 ] # Bug ? if cd_type == 'V' : ptab [..., -1, 0: ] = psgn * ptab [..., -3, -1::-1 ] ptab [..., -2, jpi//2: ] = psgn * ptab [..., -2, jpi//2-1::-1 ] if cd_type == 'F' : ptab [..., -1, 0:-1 ] = psgn * ptab [..., -3, -2::-1 ] ptab [..., -1, -1 ] = psgn * ptab [..., -3, 0 ] ptab [..., -2, jpi//2:-1] = psgn * ptab [..., -2, jpi//2-2::-1 ] return ptab def lbc_mask (ptab, nperio=None, cd_type='T', sval=np.nan) : # """ Mask fields on duplicated points ptab : Input array rank 2 at least : ptab[...., lat, lon] nperio : Type of periodicity cd_type : Grid specification : T, U, V or F See NEMO documentation for further details """ jpj, jpi, nperio, cd_type = lbc_init (ptab, nperio, cd_type) # #> East-West boundary conditions # ------------------------------ if nperio in [1, 4, 6] : # ... cyclic ptab [...,:, 0] = sval ptab [...,:, -1] = sval # #> South (in which nperio cases ?) # -------------------------------- if nperio in [1, 3, 4, 5, 6] : ptab [...,0,:] = sval # #> North-South boundary conditions # -------------------------------- if nperio in [3, 4] : # North fold T-point pivot if cd_type in [ 'T', 'W' ] : # T-, W-point ptab [..., -1, 1: ] = sval ptab [..., -1, 0 ] = sval ptab [..., -2, jpi//2: ] = sval if cd_type == 'U' : ptab [..., -1, : ] = sval ptab [..., -2, jpi//2: ] = sval if cd_type == 'V' : ptab [..., -2, : ] = sval ptab [..., -1, : ] = sval if cd_type == 'F' : ptab [..., -2, : ] = sval ptab [..., -1, : ] = sval if nperio in [5, 6] : # North fold F-point pivot if cd_type in ['T', 'W'] : ptab [..., -1, 0: ] = sval if cd_type == 'U' : ptab [..., -1, 0:-1 ] = sval ptab [..., -1, -1 ] = sval if cd_type == 'V' : ptab [..., -1, 0: ] = sval ptab [..., -2, jpi//2: ] = sval if cd_type == 'F' : ptab [..., -1, 0:-1 ] = sval ptab [..., -1, -1 ] = sval ptab [..., -2, jpi//2:-1] = sval return ptab def lbc_plot (ptab, nperio=None, cd_type='T', psgn=1.0, sval=np.nan) : """ Set periodicity on input field, adapted for plotting for any cartopy projection ptab : Input array (works rank 2 at least : ptab[...., lat, lon] nperio : Type of periodicity cd_type : Grid specification : T, U, V or F psgn : For change of sign for vector components (1 for scalars, -1 for vector components) See NEMO documentation for further details """ jpj, jpi, nperio, cd_type = lbc_init (ptab, nperio, cd_type) psgn = ptab.dtype.type (psgn) # #> East-West boundary conditions # ------------------------------ if nperio in [1, 4, 6] : # ... cyclic ptab [..., :, 0] = ptab [..., :, -2] ptab [..., :, -1] = ptab [..., :, 1] # #> North-South boundary conditions # -------------------------------- if nperio in [3, 4] : # North fold T-point pivot if cd_type in [ 'T', 'W' ] : # T-, W-point ptab [..., -1, :] = sval ptab [..., -2, :jpi//2] = sval #ptab [..., -2, :] = sval if cd_type == 'U' : ptab [..., -1, : ] = sval if cd_type == 'V' : ptab [..., -2, : ] = sval ptab [..., -1, : ] = sval if cd_type == 'F' : ptab [..., -2, : ] = sval ptab [..., -1, : ] = sval if nperio in [5, 6] : # North fold F-point pivot if cd_type in ['T', 'W'] : ptab [..., -1, : ] = sval if cd_type == 'U' : ptab [..., -1, : ] = sval if cd_type == 'V' : ptab [..., -1, : ] = sval ptab [..., -2, jpi//2: ] = sval if cd_type == 'F' : ptab [..., -1, : ] = sval ptab [..., -2, jpi//2:-1] = sval return ptab def geo2oce (pxx, pyy, pzz, glam, gphi) : ''' Change vector from geocentric to east/north Input pxx, pyy, pzz : components on the geoce,tric system glam, gphi : longitue and latitude of the points ''' gsinlon = np.sin (rad * glam) gcoslon = np.cos (rad * glam) gsinlat = np.sin (rad * gphi) gcoslat = np.cos (rad * gphi) pte = - pxx * gsinlon + pyy * gcoslon ptn = - pxx * gcoslon * gsinlat - pyy * gsinlon * gsinlat + pzz * gcoslat return pte, ptn def oce2geo (pte, ptn, glam, gphi) : ##---------------------------------------------------------------------- ## *** ROUTINE oce2geo *** ## ## ** Purpose : ## ## ** Method : Change vector from east/north to geocentric ## ## History : ## # (A. Caubel) oce2geo - Original code ##---------------------------------------------------------------------- gsinlon = np.sin (rad * glam) gcoslon = np.cos (rad * glam) gsinlat = np.sin (rad * gphi) gcoslat = np.cos (rad * gphi) pxx = - pte * gsinlon - ptn * gcoslon * gsinlat pyy = pte * gcoslon - ptn * gsinlon * gsinlat pzz = ptn * gcoslat return pxx, pyy, pzz ## =========================================================================== ## ## That's all folk's !!! ## ## ===========================================================================