Context Navigation

nemo.py @ 6709

Last change on this file since 6709 was 6665, checked in by omamce, 9 months ago

O.M. : WATER BUDGET

Improved code with pylint analysis

Property svn:keywords set to Date Revision HeadURL Author

File size: 92.7 KB

Line
1	# -- coding: utf-8 --
2	## ===========================================================================
3	##
4	## This software is governed by the CeCILL license under French law and
5	## abiding by the rules of distribution of free software. You can use,
6	## modify and/ or redistribute the software under the terms of the CeCILL
7	## license as circulated by CEA, CNRS and INRIA at the following URL
8	## "http://www.cecill.info".
9	##
10	## Warning, to install, configure, run, use any of Olivier Marti's
11	## software or to read the associated documentation you'll need at least
12	## one (1) brain in a reasonably working order. Lack of this implement
13	## will void any warranties (either express or implied).
14	## O. Marti assumes no responsability for errors, omissions,
15	## data loss, or any other consequences caused directly or indirectly by
16	## the usage of his software by incorrectly or partially configured
17	## personal.
18	##
19	## ===========================================================================
20	'''
21	Utilities to plot NEMO ORCA fields,
22	Handles periodicity and other stuff
23
24	- Lots of tests for xarray object
25	- Not much tested for numpy objects
26
27	Author: olivier.marti@lsce.ipsl.fr
28
29	## SVN information
30	__Author__ = "$Author$"
31	__Date__ = "$Date$"
32	__Revision__ = "$Revision$"
33	__Id__ = "$Id: $"
34	__HeadURL = "$HeadURL$"
35	'''
36
37	import numpy as np
38	import xarray as xr
39
40	try :
41	from sklearn.impute import SimpleImputer
42	except ImportError as err :
43	print ("Import error of sklearn.impute.SimpleImputer :", err)
44	SimpleImputer = None
45
46	try :
47	import f90nml
48	except ImportError as err :
49	print ("Import error of f90nml :", err)
50	f90nml = None
51
52
53	RPI = np.pi
54	RAD = np.deg2rad (1.0)
55	DAR = np.rad2deg (1.0)
56	REPSI = np.finfo (1.0).eps
57
58	NPERIO_VALID_RANGE = [0, 1, 4, 4.2, 5, 6, 6.2]
59
60	RAAMO = 12 # Number of months in one year
61	RJJHH = 24 # Number of hours in one day
62	RHHMM = 60 # Number of minutes in one hour
63	RMMSS = 60 # Number of seconds in one minute
64	RA = 6371229.0 # Earth radius [m]
65	GRAV = 9.80665 # Gravity [m/s2]
66	RT0 = 273.15 # Freezing point of fresh water [Kelvin]
67	RAU0 = 1026.0 # Volumic mass of sea water [kg/m3]
68	SICE = 6.0 # Salinity of ice (for pisces) [psu]
69	SOCE = 34.7 # Salinity of sea (for pisces and isf) [psu]
70	RLEVAP = 2.5e+6 # Latent heat of evaporation (water) [J/K]
71	VKARMN = 0.4 # Von Karman constant
72	STEFAN = 5.67e-8 # Stefan-Boltzmann constant [W/m2/K4]
73	RHOS = 330. # Volumic mass of snow [kg/m3]
74	RHOI = 917. # Volumic mass of sea ice [kg/m3]
75	RHOW = 1000. # Volumic mass of freshwater in melt ponds [kg/m3]
76	RCND_I = 2.034396 # Thermal conductivity of fresh ice [W/m/K]
77	RCPI = 2067.0 # Specific heat of fresh ice [J/kg/K]
78	RLSUB = 2.834e+6 # Pure ice latent heat of sublimation [J/kg]
79	RLFUS = 0.334e+6 # Latent heat of fusion of fresh ice [J/kg]
80	RTMLT = 0.054 # Decrease of seawater meltpoint with salinity
81
82	RDAY = RJJHH * RHHMM * RMMSS # Day length [s]
83	RSIYEA = 365.25 * RDAY * 2. * RPI / 6.283076 # Sideral year length [s]
84	RSIDAY = RDAY / (1. + RDAY / RSIYEA) # Sideral day length [s]
85	OMEGA = 2. * RPI / RSIDAY # Earth rotation parameter [s-1]
86
87	## Default names of dimensions
88	UDIMS = {'x':'x', 'y':'y', 'z':'olevel', 't':'time_counter'}
89
90	## All possibles name of dimensions in Nemo files
91	XNAME = [ 'x', 'X', 'X1', 'xx', 'XX',
92	'x_grid_T', 'x_grid_U', 'x_grid_V', 'x_grid_F', 'x_grid_W',
93	'lon', 'nav_lon', 'longitude', 'X1', 'x_c', 'x_f', ]
94	YNAME = [ 'y', 'Y', 'Y1', 'yy', 'YY',
95	'y_grid_T', 'y_grid_U', 'y_grid_V', 'y_grid_F', 'y_grid_W',
96	'lat', 'nav_lat', 'latitude' , 'Y1', 'y_c', 'y_f', ]
97	ZNAME = [ 'z', 'Z', 'Z1', 'zz', 'ZZ', 'depth', 'tdepth', 'udepth',
98	'vdepth', 'wdepth', 'fdepth', 'deptht', 'depthu',
99	'depthv', 'depthw', 'depthf', 'olevel', 'z_c', 'z_f', ]
100	TNAME = [ 't', 'T', 'tt', 'TT', 'time', 'time_counter', 'time_centered', ]
101
102	## All possibles name of units of dimensions in Nemo files
103	XUNIT = [ 'degrees_east', ]
104	YUNIT = [ 'degrees_north', ]
105	ZUNIT = [ 'm', 'meter', ]
106	TUNIT = [ 'second', 'minute', 'hour', 'day', 'month', 'year', ]
107
108	## All possibles size of dimensions in Orca files
109	XLENGTH = [ 180, 182, 360, 362, 1440 ]
110	YLENGTH = [ 148, 149, 331, 332 ]
111	ZLENGTH = [ 31, 75]
112
113	## ===========================================================================
114	def __mmath__ (ptab, default=None) :
115	'''Determines the type of tab : xarray, numpy or numpy.ma object ?
116
117	Returns type
118	'''
119	mmath = default
120	if isinstance (ptab, xr.core.dataarray.DataArray) :
121	mmath = xr
122	if isinstance (ptab, np.ndarray) :
123	mmath = np
124	if isinstance (ptab, np.ma.MaskType) :
125	mmath = np.ma
126
127	return mmath
128
129	def __guess_nperio__ (jpj, jpi, nperio=None, out='nperio') :
130	'''Tries to guess the value of nperio (periodicity parameter.
131
132	See NEMO documentation for details)
133	Inputs
134	jpj : number of latitudes
135	jpi : number of longitudes
136	nperio : periodicity parameter
137	'''
138	if nperio is None :
139	nperio = __guess_config__ (jpj, jpi, nperio=None, out=out)
140	return nperio
141
142	def __guess_config__ (jpj, jpi, nperio=None, config=None, out='nperio') :
143	'''Tries to guess the value of nperio (periodicity parameter).
144
145	See NEMO documentation for details)
146	Inputs
147	jpj : number of latitudes
148	jpi : number of longitudes
149	nperio : periodicity parameter
150	'''
151	print ( jpi, jpj)
152	if nperio is None :
153	## Values for NEMO version < 4.2
154	if ( (jpj == 149 and jpi == 182) or (jpj is None and jpi == 182) or
155	(jpj == 149 or jpi is None) ) :
156	# ORCA2. We choose legacy orca2.
157	config, nperio, iperio, jperio, nfold, nftype = 'ORCA2.3' , 4, 1, 0, 1, 'T'
158	if ((jpj == 332 and jpi == 362) or (jpj is None and jpi == 362) or
159	(jpj == 332 and jpi is None) ) : # eORCA1.
160	config, nperio, iperio, jperio, nfold, nftype = 'eORCA1.2', 6, 1, 0, 1, 'F'
161	if jpi == 1442 : # ORCA025.
162	config, nperio, iperio, jperio, nfold, nftype = 'ORCA025' , 6, 1, 0, 1, 'F'
163	if jpj == 294 : # ORCA1
164	config, nperio, iperio, jperio, nfold, nftype = 'ORCA1' , 6, 1, 0, 1, 'F'
165
166	## Values for NEMO version >= 4.2. No more halo points
167	if (jpj == 148 and jpi == 180) or (jpj is None and jpi == 180) or \
168	(jpj == 148 and jpi is None) : # ORCA2. We choose legacy orca2.
169	config, nperio, iperio, jperio, nfold, nftype = 'ORCA2.4' , 4.2, 1, 0, 1, 'F'
170	if (jpj == 331 and jpi == 360) or (jpj is None and jpi == 360) or \
171	(jpj == 331 and jpi is None) : # eORCA1.
172	config, nperio, iperio, jperio, nfold, nftype = 'eORCA1.4', 6.2, 1, 0, 1, 'F'
173	if jpi == 1440 : # ORCA025.
174	config, nperio, iperio, jperio, nfold, nftype = 'ORCA025' , 6.2, 1, 0, 1, 'F'
175
176	if nperio is None :
177	raise ValueError ('in nemo module : nperio not found, and cannot by guessed')
178
179	if nperio in NPERIO_VALID_RANGE :
180	print ( f'nperio set as {nperio} (deduced from {jpj=} and {jpi=})' )
181	else :
182	raise ValueError ( f'nperio set as {nperio} (deduced from {jpi=} and {jpj=}) : \n'+
183	'nemo.py is not ready for this value' )
184
185	if out == 'nperio' :
186	return nperio
187	if out == 'config' :
188	return config
189	if out == 'perio' :
190	return iperio, jperio, nfold, nftype
191	if out in ['full', 'all'] :
192	return {'nperio':nperio, 'iperio':iperio, 'jperio':jperio, 'nfold':nfold, 'nftype':nftype}
193
194	def __guess_point__ (ptab) :
195	'''Tries to guess the grid point (periodicity parameter.
196
197	See NEMO documentation for details)
198	For array conforments with xgcm requirements
199
200	Inputs
201	ptab : xarray array
202
203	Credits : who is the original author ?
204	'''
205
206	gp = None
207	mmath = __mmath__ (ptab)
208	if mmath == xr :
209	if ('x_c' in ptab.dims and 'y_c' in ptab.dims ) :
210	gp = 'T'
211	if ('x_f' in ptab.dims and 'y_c' in ptab.dims ) :
212	gp = 'U'
213	if ('x_c' in ptab.dims and 'y_f' in ptab.dims ) :
214	gp = 'V'
215	if ('x_f' in ptab.dims and 'y_f' in ptab.dims ) :
216	gp = 'F'
217	if ('x_c' in ptab.dims and 'y_c' in ptab.dims
218	and 'z_c' in ptab.dims ) :
219	gp = 'T'
220	if ('x_c' in ptab.dims and 'y_c' in ptab.dims
221	and 'z_f' in ptab.dims ) :
222	gp = 'W'
223	if ('x_f' in ptab.dims and 'y_c' in ptab.dims
224	and 'z_f' in ptab.dims ) :
225	gp = 'U'
226	if ('x_c' in ptab.dims and 'y_f' in ptab.dims
227	and 'z_f' in ptab.dims ) :
228	gp = 'V'
229	if ('x_f' in ptab.dims and 'y_f' in ptab.dims
230	and 'z_f' in ptab.dims ) :
231	gp = 'F'
232
233	if gp is None :
234	raise AttributeError ('in nemo module : cd_type not found, and cannot by guessed')
235	print ( f'Grid set as {gp} deduced from dims {ptab.dims}' )
236	return gp
237	else :
238	raise AttributeError ('in nemo module : cd_type not found, input is not an xarray data')
239
240	def get_shape ( ptab ) :
241	'''Get shape of ptab return a string with axes names
242
243	shape may contain X, Y, Z or T
244	Y is missing for a latitudinal slice
245	X is missing for on longitudinal slice
246	etc ...
247	'''
248
249	g_shape = ''
250	if __find_axis__ (ptab, 'x')[0] :
251	g_shape = 'X'
252	if __find_axis__ (ptab, 'y')[0] :
253	g_shape = 'Y' + g_shape
254	if __find_axis__ (ptab, 'z')[0] :
255	g_shape = 'Z' + g_shape
256	if __find_axis__ (ptab, 't')[0] :
257	g_shape = 'T' + g_shape
258	return g_shape
259
260	def lbc_diag (nperio) :
261	'''Useful to switch between field with and without halo'''
262	lperio, aperio = nperio, False
263	if nperio == 4.2 :
264	lperio, aperio = 4, True
265	if nperio == 6.2 :
266	lperio, aperio = 6, True
267	return lperio, aperio
268
269	def __find_axis__ (ptab, axis='z', back=True) :
270	'''Returns name and name of the requested axis'''
271	mmath = __mmath__ (ptab)
272	ax, ix = None, None
273
274	if axis in XNAME :
275	ax_name, unit_list, length = XNAME, XUNIT, XLENGTH
276	if axis in YNAME :
277	ax_name, unit_list, length = YNAME, YUNIT, YLENGTH
278	if axis in ZNAME :
279	ax_name, unit_list, length = ZNAME, ZUNIT, ZLENGTH
280	if axis in TNAME :
281	ax_name, unit_list, length = TNAME, TUNIT, None
282
283	if mmath == xr :
284	# Try by name
285	for dim in ax_name :
286	if dim in ptab.dims :
287	ix, ax = ptab.dims.index (dim), dim
288
289	# If not found, try by axis attributes
290	if not ix :
291	for i, dim in enumerate (ptab.dims) :
292	if 'axis' in ptab.coords[dim].attrs.keys() :
293	l_axis = ptab.coords[dim].attrs['axis']
294	if axis in ax_name and l_axis == 'X' :
295	ix, ax = (i, dim)
296	if axis in ax_name and l_axis == 'Y' :
297	ix, ax = (i, dim)
298	if axis in ax_name and l_axis == 'Z' :
299	ix, ax = (i, dim)
300	if axis in ax_name and l_axis == 'T' :
301	ix, ax = (i, dim)
302
303	# If not found, try by units
304	if not ix :
305	for i, dim in enumerate (ptab.dims) :
306	if 'units' in ptab.coords[dim].attrs.keys() :
307	for name in unit_list :
308	if name in ptab.coords[dim].attrs['units'] :
309	ix, ax = i, dim
310
311	# If numpy array or dimension not found, try by length
312	if mmath != xr or not ix :
313	if length :
314	l_shape = ptab.shape
315	for nn in np.arange ( len(l_shape) ) :
316	if l_shape[nn] in length :
317	ix = nn
318
319	if ix and back :
320	ix -= len(ptab.shape)
321
322	return ax, ix
323
324	def find_axis ( ptab, axis='z', back=True ) :
325	'''Version of find_axis with no __'''
326	ix, xx = __find_axis__ (ptab, axis, back)
327	return xx, ix
328
329	def fixed_lon (plon, center_lon=0.0) :
330	'''Returns corrected longitudes for nicer plots
331
332	lon : longitudes of the grid. At least 2D.
333	center_lon : center longitude. Default=0.
334
335	Designed by Phil Pelson.
336	See https://gist.github.com/pelson/79cf31ef324774c97ae7
337	'''
338	mmath = __mmath__ (plon)
339
340	f_lon = plon.copy ()
341
342	f_lon = mmath.where (f_lon > center_lon+180., f_lon-360.0, f_lon)
343	f_lon = mmath.where (f_lon < center_lon-180., f_lon+360.0, f_lon)
344
345	for i, start in enumerate (np.argmax (np.abs (np.diff (f_lon, axis=-1)) > 180., axis=-1)) :
346	f_lon [..., i, start+1:] += 360.
347
348	# Special case for eORCA025
349	if f_lon.shape [-1] == 1442 :
350	f_lon [..., -2, :] = f_lon [..., -3, :]
351	if f_lon.shape [-1] == 1440 :
352	f_lon [..., -1, :] = f_lon [..., -2, :]
353
354	if f_lon.min () > center_lon :
355	f_lon += -360.0
356	if f_lon.max () < center_lon :
357	f_lon += 360.0
358
359	if f_lon.min () < center_lon-360.0 :
360	f_lon += 360.0
361	if f_lon.max () > center_lon+360.0 :
362	f_lon += -360.0
363
364	return f_lon
365
366	def bounds_clolon ( pbounds_lon, plon, rad=False, deg=True) :
367	'''Choose closest to lon0 longitude, adding/substacting 360Â° if needed
368	'''
369
370	if rad :
371	lon_range = 2.0*np.pi
372	if deg :
373	lon_range = 360.0
374	b_clolon = pbounds_lon.copy ()
375
376	b_clolon = xr.where ( b_clolon < plon-lon_range/2.,
377	b_clolon+lon_range,
378	b_clolon )
379	b_clolon = xr.where ( b_clolon > plon+lon_range/2.,
380	b_clolon-lon_range,
381	b_clolon )
382	return b_clolon
383
384	def unify_dims ( dd, x='x', y='y', z='olevel', t='time_counter', verbose=False ) :
385	'''Rename dimensions to unify them between NEMO versions
386	'''
387	for xx in XNAME :
388	if xx in dd.dims and xx != x :
389	if verbose :
390	print ( f"{xx} renamed to {x}" )
391	dd = dd.rename ( {xx:x})
392
393	for yy in YNAME :
394	if yy in dd.dims and yy != y :
395	if verbose :
396	print ( f"{yy} renamed to {y}" )
397	dd = dd.rename ( {yy:y} )
398
399	for zz in ZNAME :
400	if zz in dd.dims and zz != z :
401	if verbose :
402	print ( f"{zz} renamed to {z}" )
403	dd = dd.rename ( {zz:z} )
404
405	for tt in TNAME :
406	if tt in dd.dims and tt != t :
407	if verbose :
408	print ( f"{tt} renamed to {t}" )
409	dd = dd.rename ( {tt:t} )
410
411	return dd
412
413
414	if SimpleImputer :
415	def fill_empty (ptab, sval=np.nan, transpose=False) :
416	'''Fill empty values
417
418	Useful when NEMO has run with no wet points options :
419	some parts of the domain, with no ocean points, have no
420	values
421	'''
422	mmath = __mmath__ (ptab)
423
424	imp = SimpleImputer (missing_values=sval, strategy='mean')
425	if transpose :
426	imp.fit (ptab.T)
427	ztab = imp.transform (ptab.T).T
428	else :
429	imp.fit (ptab)
430	ztab = imp.transform (ptab)
431
432	if mmath == xr :
433	ztab = xr.DataArray (ztab, dims=ztab.dims, coords=ztab.coords)
434	ztab.attrs.update (ptab.attrs)
435
436	return ztab
437
438
439	else :
440	print ("Import error of sklearn.impute.SimpleImputer")
441	def fill_empty (ptab, sval=np.nan, transpose=False) :
442	'''Void version of fill_empy, because module sklearn.impute.SimpleImputer is not available
443
444	fill_empty :
445	Fill values
446
447	Useful when NEMO has run with no wet points options :
448	some parts of the domain, with no ocean points, have no
449	values
450	'''
451	print ( 'Error : module sklearn.impute.SimpleImputer not found' )
452	print ( 'Can not call fill_empty' )
453	print ( 'Call arguments where : ' )
454	print ( f'{ptab.shape=} {sval=} {transpose=}' )
455
456	def fill_lonlat (plon, plat, sval=-1) :
457	'''Fill longitude/latitude values
458
459	Useful when NEMO has run with no wet points options :
460	some parts of the domain, with no ocean points, have no
461	lon/lat values
462	'''
463	from sklearn.impute import SimpleImputer
464	mmath = __mmath__ (plon)
465
466	imp = SimpleImputer (missing_values=sval, strategy='mean')
467	imp.fit (plon)
468	zlon = imp.transform (plon)
469	imp.fit (plat.T)
470	zlat = imp.transform (plat.T).T
471
472	if mmath == xr :
473	zlon = xr.DataArray (zlon, dims=plon.dims, coords=plon.coords)
474	zlat = xr.DataArray (zlat, dims=plat.dims, coords=plat.coords)
475	zlon.attrs.update (plon.attrs)
476	zlat.attrs.update (plat.attrs)
477
478	zlon = fixed_lon (zlon)
479
480	return zlon, zlat
481
482	def fill_bounds_lonlat (pbounds_lon, pbounds_lat, sval=-1) :
483	'''Fill longitude/latitude bounds values
484
485	Useful when NEMO has run with no wet points options :
486	some parts of the domain, with no ocean points, as no
487	lon/lat values
488	'''
489	mmath = __mmath__ (pbounds_lon)
490
491	z_bounds_lon = np.empty ( pbounds_lon.shape )
492	z_bounds_lat = np.empty ( pbounds_lat.shape )
493
494	imp = SimpleImputer (missing_values=sval, strategy='mean')
495
496	for n in np.arange (4) :
497	imp.fit (pbounds_lon[:,:,n])
498	z_bounds_lon[:,:,n] = imp.transform (pbounds_lon[:,:,n])
499	imp.fit (pbounds_lat[:,:,n].T)
500	z_bounds_lat[:,:,n] = imp.transform (pbounds_lat[:,:,n].T).T
501
502	if mmath == xr :
503	z_bounds_lon = xr.DataArray (pbounds_lon, dims=pbounds_lon.dims,
504	coords=pbounds_lon.coords)
505	z_bounds_lat = xr.DataArray (pbounds_lat, dims=pbounds_lat.dims,
506	coords=pbounds_lat.coords)
507	z_bounds_lon.attrs.update (pbounds_lat.attrs)
508	z_bounds_lat.attrs.update (pbounds_lat.attrs)
509
510	return z_bounds_lon, z_bounds_lat
511
512	def jeq (plat) :
513	'''Returns j index of equator in the grid
514
515	lat : latitudes of the grid. At least 2D.
516	'''
517	mmath = __mmath__ (plat)
518	jy = __find_axis__ (plat, 'y')[-1]
519
520	if mmath == xr :
521	jj = int ( np.mean ( np.argmin (np.abs (np.float64 (plat)),
522	axis=jy) ) )
523	else :
524	jj = np.argmin (np.abs (np.float64 (plat[...,:, 0])))
525
526	return jj
527
528	def lon1d (plon, plat=None) :
529	'''Returns 1D longitude for simple plots.
530
531	plon : longitudes of the grid
532	plat (optionnal) : latitudes of the grid
533	'''
534	mmath = __mmath__ (plon)
535	jpj, jpi = plon.shape [-2:]
536	if np.max (plat) :
537	je = jeq (plat)
538	lon0 = plon [..., je, 0].copy()
539	dlon = plon [..., je, 1].copy() - plon [..., je, 0].copy()
540	lon_1d = np.linspace ( start=lon0, stop=lon0+360.+2*dlon, num=jpi )
541	else :
542	lon0 = plon [..., jpj//3, 0].copy()
543	dlon = plon [..., jpj//3, 1].copy() - plon [..., jpj//3, 0].copy()
544	lon_1d = np.linspace ( start=lon0, stop=lon0+360.+2*dlon, num=jpi )
545
546	#start = np.argmax (np.abs (np.diff (lon1D, axis=-1)) > 180.0, axis=-1)
547	#lon1D [..., start+1:] += 360
548
549	if mmath == xr :
550	lon_1d = xr.DataArray( lon_1d, dims=('lon',), coords=(lon_1d,))
551	lon_1d.attrs.update (plon.attrs)
552	lon_1d.attrs['units'] = 'degrees_east'
553	lon_1d.attrs['standard_name'] = 'longitude'
554	lon_1d.attrs['long_name :'] = 'Longitude'
555
556	return lon_1d
557
558	def latreg (plat, diff=0.1) :
559	'''Returns maximum j index where gridlines are along latitudes
560	in the northern hemisphere
561
562	lat : latitudes of the grid (2D)
563	diff [optional] : tolerance
564	'''
565	#mmath = __mmath__ (plat)
566	if diff is None :
567	dy = np.float64 (np.mean (np.abs (plat -
568	np.roll (plat,shift=1,axis=-2, roll_coords=False))))
569	print ( f'{dy=}' )
570	diff = dy/100.
571
572	je = jeq (plat)
573	jreg = np.where (plat[...,je:,:].max(axis=-1) -
574	plat[...,je:,:].min(axis=-1)< diff)[-1][-1] + je
575	lareg = np.float64 (plat[...,jreg,:].mean(axis=-1))
576
577	return jreg, lareg
578
579	def lat1d (plat) :
580	'''Returns 1D latitudes for zonal means and simple plots.
581
582	plat : latitudes of the grid (2D)
583	'''
584	mmath = __mmath__ (plat)
585	iy = __find_axis__ (plat, 'y')[-1]
586	jpj = plat.shape[iy]
587
588	dy = np.float64 (np.mean (np.abs (plat - np.roll (plat, shift=1,axis=-2))))
589	je = jeq (plat)
590	lat_eq = np.float64 (plat[...,je,:].mean(axis=-1))
591
592	jreg, lat_reg = latreg (plat)
593	lat_ave = np.mean (plat, axis=-1)
594
595	if np.abs (lat_eq) < dy/100. : # T, U or W grid
596	if jpj-1 > jreg :
597	dys = (90.-lat_reg) / (jpj-jreg-1)*0.5
598	else :
599	dys = (90.-lat_reg) / 2.0
600	yrange = 90.-dys-lat_reg
601	else : # V or F grid
602	yrange = 90.-lat_reg
603
604	if jpj-1 > jreg :
605	lat_1d = mmath.where (lat_ave<lat_reg,
606	lat_ave,
607	lat_reg + yrange * (np.arange(jpj)-jreg)/(jpj-jreg-1) )
608	else :
609	lat_1d = lat_ave
610	lat_1d[-1] = 90.0
611
612	if mmath == xr :
613	lat_1d = xr.DataArray( lat_1d.values, dims=('lat',), coords=(lat_1d,))
614	lat_1d.attrs.update (plat.attrs)
615	lat_1d.attrs ['units'] = 'degrees_north'
616	lat_1d.attrs ['standard_name'] = 'latitude'
617	lat_1d.attrs ['long_name :'] = 'Latitude'
618
619	return lat_1d
620
621	def latlon1d (plat, plon) :
622	'''Returns simple latitude and longitude (1D) for simple plots.
623
624	plat, plon : latitudes and longitudes of the grid (2D)
625	'''
626	return lat1d (plat), lon1d (plon, plat)
627
628	def ff (plat) :
629	'''Returns Coriolis factor
630	'''
631	zff = np.sin (RAD * plat) * OMEGA
632	return zff
633
634	def beta (plat) :
635	'''Return Beta factor (derivative of Coriolis factor)
636	'''
637	zbeta = np.cos (RAD * plat) * OMEGA / RA
638	return zbeta
639
640	def mask_lonlat (ptab, x0, x1, y0, y1, lon, lat, sval=np.nan) :
641	'''Returns masked values outside a lat/lon box
642	'''
643	mmath = __mmath__ (ptab)
644	if mmath == xr :
645	lon = lon.copy().to_masked_array()
646	lat = lat.copy().to_masked_array()
647
648	mask = np.logical_and (np.logical_and(lat>y0, lat<y1),
649	np.logical_or (np.logical_or (
650	np.logical_and(lon>x0, lon<x1),
651	np.logical_and(lon+360>x0, lon+360<x1)),
652	np.logical_and(lon-360>x0, lon-360<x1)))
653	tab = mmath.where (mask, ptab, sval)
654
655	return tab
656
657	def extend (ptab, blon=False, jplus=25, jpi=None, nperio=4) :
658	'''Returns extended field eastward to have better plots,
659	and box average crossing the boundary
660
661	Works only for xarray and numpy data (?)
662	Useful for plotting vertical sections in OCE and ATM.
663
664	ptab : field to extend.
665	blon : (optional, default=False) : if True, add 360 in the extended
666	parts of the field
667	jpi : normal longitude dimension of the field. extend does nothing
668	if the actual size of the field != jpi
669	(avoid to extend several times in notebooks)
670	jplus (optional, default=25) : number of points added on
671	the east side of the field
672
673	'''
674	mmath = __mmath__ (ptab)
675
676	if ptab.shape[-1] == 1 :
677	tabex = ptab
678
679	else :
680	if jpi is None :
681	jpi = ptab.shape[-1]
682
683	if blon :
684	xplus = -360.0
685	else :
686	xplus = 0.0
687
688	if ptab.shape[-1] > jpi :
689	tabex = ptab
690	else :
691	if nperio in [ 0, 4.2 ] :
692	istart, le, la = 0, jpi+1, 0
693	if nperio == 1 :
694	istart, le, la = 0, jpi+1, 0
695	if nperio in [4, 6] : # OPA case with two halo points for periodicity
696	# Perfect, except at the pole that should be masked by lbc_plot
697	istart, le, la = 1, jpi-2, 1
698	if mmath == xr :
699	tabex = np.concatenate (
700	(ptab.values[..., istart :istart+le+1 ] + xplus,
701	ptab.values[..., istart+la:istart+la+jplus] ),
702	axis=-1)
703	lon = ptab.dims[-1]
704	new_coords = []
705	for coord in ptab.dims :
706	if coord == lon :
707	new_coords.append ( np.arange( tabex.shape[-1]))
708	else :
709	new_coords.append ( ptab.coords[coord].values)
710	tabex = xr.DataArray ( tabex, dims=ptab.dims,
711	coords=new_coords )
712	else :
713	tabex = np.concatenate (
714	(ptab [..., istart :istart+le+1 ] + xplus,
715	ptab [..., istart+la:istart+la+jplus] ),
716	axis=-1)
717	return tabex
718
719	def orca2reg (dd, lat_name='nav_lat', lon_name='nav_lon',
720	y_name='y', x_name='x') :
721	'''Assign an ORCA dataset on a regular grid.
722
723	For use in the tropical region.
724	Inputs :
725	ff : xarray dataset
726	lat_name, lon_name : name of latitude and longitude 2D field in ff
727	y_name, x_name : namex of dimensions in ff
728
729	Returns : xarray dataset with rectangular grid. Incorrect above 20Â°N
730	'''
731	# Compute 1D longitude and latitude
732	(zlat, zlon) = latlon1d ( dd[lat_name], dd[lon_name])
733
734	zdd = dd
735	# Assign lon and lat as dimensions of the dataset
736	if y_name in zdd.dims :
737	zlat = xr.DataArray (zlat, coords=[zlat,], dims=['lat',])
738	zdd = zdd.rename_dims ({y_name: "lat",}).assign_coords (lat=zlat)
739	if x_name in zdd.dims :
740	zlon = xr.DataArray (zlon, coords=[zlon,], dims=['lon',])
741	zdd = zdd.rename_dims ({x_name: "lon",}).assign_coords (lon=zlon)
742	# Force dimensions to be in the right order
743	coord_order = ['lat', 'lon']
744	for dim in [ 'depthw', 'depthv', 'depthu', 'deptht', 'depth', 'z',
745	'time_counter', 'time', 'tbnds',
746	'bnds', 'axis_nbounds', 'two2', 'two1', 'two', 'four',] :
747	if dim in zdd.dims :
748	coord_order.insert (0, dim)
749
750	zdd = zdd.transpose (*coord_order)
751	return zdd
752
753	def lbc_init (ptab, nperio=None) :
754	'''Prepare for all lbc calls
755
756	Set periodicity on input field
757	nperio : Type of periodicity
758	0 : No periodicity
759	1, 4, 6 : Cyclic on i dimension (generaly longitudes) with 2 points halo
760	2 : Obsolete (was symmetric condition at southern boundary ?)
761	3, 4 : North fold T-point pivot (legacy ORCA2)
762	5, 6 : North fold F-point pivot (ORCA1, ORCA025, ORCA2 with new grid for paleo)
763	cd_type : Grid specification : T, U, V or F
764
765	See NEMO documentation for further details
766	'''
767	jpi, jpj = None, None
768	ax, ix = __find_axis__ (ptab, 'x')
769	ay, jy = __find_axis__ (ptab, 'y')
770	if ax :
771	jpi = ptab.shape[ix]
772	if ay :
773	jpj = ptab.shape[jy]
774
775	if nperio is None :
776	nperio = __guess_nperio__ (jpj, jpi, nperio)
777
778	if nperio not in NPERIO_VALID_RANGE :
779	raise AttributeError ( f'{nperio=} is not in the valid range {NPERIO_VALID_RANGE}' )
780
781	return jpj, jpi, nperio
782
783	def lbc (ptab, nperio=None, cd_type='T', psgn=1.0, nemo_4u_bug=False) :
784	'''Set periodicity on input field
785
786	ptab : Input array (works for rank 2 at least : ptab[...., lat, lon])
787	nperio : Type of periodicity
788	cd_type : Grid specification : T, U, V or F
789	psgn : For change of sign for vector components (1 for scalars, -1 for vector components)
790
791	See NEMO documentation for further details
792	'''
793	jpi, nperio = lbc_init (ptab, nperio)[1:]
794	ax = __find_axis__ (ptab, 'x')[0]
795	ay = __find_axis__ (ptab, 'y')[0]
796	psgn = ptab.dtype.type (psgn)
797	mmath = __mmath__ (ptab)
798
799	if mmath == xr :
800	ztab = ptab.values.copy ()
801	else :
802	ztab = ptab.copy ()
803
804	if ax :
805	#
806	#> East-West boundary conditions
807	# ------------------------------
808	if nperio in [1, 4, 6] :
809	# ... cyclic
810	ztab [..., 0] = ztab [..., -2]
811	ztab [..., -1] = ztab [..., 1]
812
813	if ay :
814	#
815	#> North-South boundary conditions
816	# --------------------------------
817	if nperio in [3, 4] : # North fold T-point pivot
818	if cd_type in [ 'T', 'W' ] : # T-, W-point
819	ztab [..., -1, 1: ] = psgn * ztab [..., -3, -1:0:-1 ]
820	ztab [..., -1, 0 ] = psgn * ztab [..., -3, 2 ]
821	ztab [..., -2, jpi//2: ] = psgn * ztab [..., -2, jpi//2:0:-1 ]
822
823	if cd_type == 'U' :
824	ztab [..., -1, 0:-1 ] = psgn * ztab [..., -3, -1:0:-1 ]
825	ztab [..., -1, 0 ] = psgn * ztab [..., -3, 1 ]
826	ztab [..., -1, -1 ] = psgn * ztab [..., -3, -2 ]
827
828	if nemo_4u_bug :
829	ztab [..., -2, jpi//2+1:-1] = psgn * ztab [..., -2, jpi//2-2:0:-1]
830	ztab [..., -2, jpi//2-1 ] = psgn * ztab [..., -2, jpi//2 ]
831	else :
832	ztab [..., -2, jpi//2-1:-1] = psgn * ztab [..., -2, jpi//2:0:-1]
833
834	if cd_type == 'V' :
835	ztab [..., -2, 1: ] = psgn * ztab [..., -3, jpi-1:0:-1 ]
836	ztab [..., -1, 1: ] = psgn * ztab [..., -4, -1:0:-1 ]
837	ztab [..., -1, 0 ] = psgn * ztab [..., -4, 2 ]
838
839	if cd_type == 'F' :
840	ztab [..., -2, 0:-1 ] = psgn * ztab [..., -3, -1:0:-1 ]
841	ztab [..., -1, 0:-1 ] = psgn * ztab [..., -4, -1:0:-1 ]
842	ztab [..., -1, 0 ] = psgn * ztab [..., -4, 1 ]
843	ztab [..., -1, -1 ] = psgn * ztab [..., -4, -2 ]
844
845	if nperio in [4.2] : # North fold T-point pivot
846	if cd_type in [ 'T', 'W' ] : # T-, W-point
847	ztab [..., -1, jpi//2: ] = psgn * ztab [..., -1, jpi//2:0:-1 ]
848
849	if cd_type == 'U' :
850	ztab [..., -1, jpi//2-1:-1] = psgn * ztab [..., -1, jpi//2:0:-1]
851
852	if cd_type == 'V' :
853	ztab [..., -1, 1: ] = psgn * ztab [..., -2, jpi-1:0:-1 ]
854
855	if cd_type == 'F' :
856	ztab [..., -1, 0:-1 ] = psgn * ztab [..., -2, -1:0:-1 ]
857
858	if nperio in [5, 6] : # North fold F-point pivot
859	if cd_type in ['T', 'W'] :
860	ztab [..., -1, 0: ] = psgn * ztab [..., -2, -1::-1 ]
861
862	if cd_type == 'U' :
863	ztab [..., -1, 0:-1 ] = psgn * ztab [..., -2, -2::-1 ]
864	ztab [..., -1, -1 ] = psgn * ztab [..., -2, 0 ] # Bug ?
865
866	if cd_type == 'V' :
867	ztab [..., -1, 0: ] = psgn * ztab [..., -3, -1::-1 ]
868	ztab [..., -2, jpi//2: ] = psgn * ztab [..., -2, jpi//2-1::-1 ]
869
870	if cd_type == 'F' :
871	ztab [..., -1, 0:-1 ] = psgn * ztab [..., -3, -2::-1 ]
872	ztab [..., -1, -1 ] = psgn * ztab [..., -3, 0 ]
873	ztab [..., -2, jpi//2:-1] = psgn * ztab [..., -2, jpi//2-2::-1 ]
874
875	#
876	#> East-West boundary conditions
877	# ------------------------------
878	if nperio in [1, 4, 6] :
879	# ... cyclic
880	ztab [..., 0] = ztab [..., -2]
881	ztab [..., -1] = ztab [..., 1]
882
883	if mmath == xr :
884	ztab = xr.DataArray ( ztab, dims=ptab.dims, coords=ptab.coords )
885	ztab.attrs = ptab.attrs
886
887	return ztab
888
889	def lbc_mask (ptab, nperio=None, cd_type='T', sval=np.nan) :
890	'''Mask fields on duplicated points
891
892	ptab : Input array. Rank 2 at least : ptab [...., lat, lon]
893	nperio : Type of periodicity
894	cd_type : Grid specification : T, U, V or F
895
896	See NEMO documentation for further details
897	'''
898	jpi, nperio = lbc_init (ptab, nperio)[1:]
899	ax = __find_axis__ (ptab, 'x')[0]
900	ay = __find_axis__ (ptab, 'y')[0]
901	ztab = ptab.copy ()
902
903	if ax :
904	#
905	#> East-West boundary conditions
906	# ------------------------------
907	if nperio in [1, 4, 6] :
908	# ... cyclic
909	ztab [..., 0] = sval
910	ztab [..., -1] = sval
911
912	if ay :
913	#
914	#> South (in which nperio cases ?)
915	# --------------------------------
916	if nperio in [1, 3, 4, 5, 6] :
917	ztab [..., 0, :] = sval
918
919	#
920	#> North-South boundary conditions
921	# --------------------------------
922	if nperio in [3, 4] : # North fold T-point pivot
923	if cd_type in [ 'T', 'W' ] : # T-, W-point
924	ztab [..., -1, : ] = sval
925	ztab [..., -2, :jpi//2 ] = sval
926
927	if cd_type == 'U' :
928	ztab [..., -1, : ] = sval
929	ztab [..., -2, jpi//2+1: ] = sval
930
931	if cd_type == 'V' :
932	ztab [..., -2, : ] = sval
933	ztab [..., -1, : ] = sval
934
935	if cd_type == 'F' :
936	ztab [..., -2, : ] = sval
937	ztab [..., -1, : ] = sval
938
939	if nperio in [4.2] : # North fold T-point pivot
940	if cd_type in [ 'T', 'W' ] : # T-, W-point
941	ztab [..., -1, jpi//2 : ] = sval
942
943	if cd_type == 'U' :
944	ztab [..., -1, jpi//2-1:-1] = sval
945
946	if cd_type == 'V' :
947	ztab [..., -1, 1: ] = sval
948
949	if cd_type == 'F' :
950	ztab [..., -1, 0:-1 ] = sval
951
952	if nperio in [5, 6] : # North fold F-point pivot
953	if cd_type in ['T', 'W'] :
954	ztab [..., -1, 0: ] = sval
955
956	if cd_type == 'U' :
957	ztab [..., -1, 0:-1 ] = sval
958	ztab [..., -1, -1 ] = sval
959
960	if cd_type == 'V' :
961	ztab [..., -1, 0: ] = sval
962	ztab [..., -2, jpi//2: ] = sval
963
964	if cd_type == 'F' :
965	ztab [..., -1, 0:-1 ] = sval
966	ztab [..., -1, -1 ] = sval
967	ztab [..., -2, jpi//2+1:-1] = sval
968
969	return ztab
970
971	def lbc_plot (ptab, nperio=None, cd_type='T', psgn=1.0, sval=np.nan) :
972	'''Set periodicity on input field, for plotting for any cartopy projection
973
974	Points at the north fold are masked
975	Points for zonal periodicity are kept
976	ptab : Input array. Rank 2 at least : ptab[...., lat, lon]
977	nperio : Type of periodicity
978	cd_type : Grid specification : T, U, V or F
979	psgn : For change of sign for vector components
980	(1 for scalars, -1 for vector components)
981
982	See NEMO documentation for further details
983	'''
984	jpi, nperio = lbc_init (ptab, nperio)[1:]
985	ax = __find_axis__ (ptab, 'x')[0]
986	ay = __find_axis__ (ptab, 'y')[0]
987	psgn = ptab.dtype.type (psgn)
988	ztab = ptab.copy ()
989
990	if ax :
991	#
992	#> East-West boundary conditions
993	# ------------------------------
994	if nperio in [1, 4, 6] :
995	# ... cyclic
996	ztab [..., :, 0] = ztab [..., :, -2]
997	ztab [..., :, -1] = ztab [..., :, 1]
998
999	if ay :
1000	#> Masks south
1001	# ------------
1002	if nperio in [4, 6] :
1003	ztab [..., 0, : ] = sval
1004
1005	#
1006	#> North-South boundary conditions
1007	# --------------------------------
1008	if nperio in [3, 4] : # North fold T-point pivot
1009	if cd_type in [ 'T', 'W' ] : # T-, W-point
1010	ztab [..., -1, : ] = sval
1011	#ztab [..., -2, jpi//2: ] = sval
1012	ztab [..., -2, :jpi//2 ] = sval # Give better plots than above
1013	if cd_type == 'U' :
1014	ztab [..., -1, : ] = sval
1015
1016	if cd_type == 'V' :
1017	ztab [..., -2, : ] = sval
1018	ztab [..., -1, : ] = sval
1019
1020	if cd_type == 'F' :
1021	ztab [..., -2, : ] = sval
1022	ztab [..., -1, : ] = sval
1023
1024	if nperio in [4.2] : # North fold T-point pivot
1025	if cd_type in [ 'T', 'W' ] : # T-, W-point
1026	ztab [..., -1, jpi//2: ] = sval
1027
1028	if cd_type == 'U' :
1029	ztab [..., -1, jpi//2-1:-1] = sval
1030
1031	if cd_type == 'V' :
1032	ztab [..., -1, 1: ] = sval
1033
1034	if cd_type == 'F' :
1035	ztab [..., -1, 0:-1 ] = sval
1036
1037	if nperio in [5, 6] : # North fold F-point pivot
1038	if cd_type in ['T', 'W'] :
1039	ztab [..., -1, : ] = sval
1040
1041	if cd_type == 'U' :
1042	ztab [..., -1, : ] = sval
1043
1044	if cd_type == 'V' :
1045	ztab [..., -1, : ] = sval
1046	ztab [..., -2, jpi//2: ] = sval
1047
1048	if cd_type == 'F' :
1049	ztab [..., -1, : ] = sval
1050	ztab [..., -2, jpi//2+1:-1] = sval
1051
1052	return ztab
1053
1054	def lbc_add (ptab, nperio=None, cd_type=None, psgn=1) :
1055	'''Handles NEMO domain changes between NEMO 4.0 to NEMO 4.2
1056
1057	Periodicity and north fold halos has been removed in NEMO 4.2
1058	This routine adds the halos if needed
1059
1060	ptab : Input array (works
1061	rank 2 at least : ptab[...., lat, lon]
1062	nperio : Type of periodicity
1063
1064	See NEMO documentation for further details
1065	'''
1066	mmath = __mmath__ (ptab)
1067	nperio = lbc_init (ptab, nperio)[-1]
1068	lshape = get_shape (ptab)
1069	ix = __find_axis__ (ptab, 'x')[-1]
1070	jy = __find_axis__ (ptab, 'y')[-1]
1071
1072	t_shape = np.array (ptab.shape)
1073
1074	if nperio in [4.2, 6.2] :
1075
1076	ext_shape = t_shape.copy()
1077	if 'X' in lshape :
1078	ext_shape[ix] = ext_shape[ix] + 2
1079	if 'Y' in lshape :
1080	ext_shape[jy] = ext_shape[jy] + 1
1081
1082	if mmath == xr :
1083	ptab_ext = xr.DataArray (np.zeros (ext_shape), dims=ptab.dims)
1084	if 'X' in lshape and 'Y' in lshape :
1085	ptab_ext.values[..., :-1, 1:-1] = ptab.values.copy ()
1086	else :
1087	if 'X' in lshape :
1088	ptab_ext.values[..., 1:-1] = ptab.values.copy ()
1089	if 'Y' in lshape :
1090	ptab_ext.values[..., :-1 ] = ptab.values.copy ()
1091	else :
1092	ptab_ext = np.zeros (ext_shape)
1093	if 'X' in lshape and 'Y' in lshape :
1094	ptab_ext [..., :-1, 1:-1] = ptab.copy ()
1095	else :
1096	if 'X' in lshape :
1097	ptab_ext [..., 1:-1] = ptab.copy ()
1098	if 'Y' in lshape :
1099	ptab_ext [..., :-1 ] = ptab.copy ()
1100
1101	if nperio == 4.2 :
1102	ptab_ext = lbc (ptab_ext, nperio=4, cd_type=cd_type, psgn=psgn)
1103	if nperio == 6.2 :
1104	ptab_ext = lbc (ptab_ext, nperio=6, cd_type=cd_type, psgn=psgn)
1105
1106	if mmath == xr :
1107	ptab_ext.attrs = ptab.attrs
1108	az = __find_axis__ (ptab, 'z')[0]
1109	at = __find_axis__ (ptab, 't')[0]
1110	if az :
1111	ptab_ext = ptab_ext.assign_coords ( {az:ptab.coords[az]} )
1112	if at :
1113	ptab_ext = ptab_ext.assign_coords ( {at:ptab.coords[at]} )
1114
1115	else : ptab_ext = lbc (ptab, nperio=nperio, cd_type=cd_type, psgn=psgn)
1116
1117	return ptab_ext
1118
1119	def lbc_del (ptab, nperio=None, cd_type='T', psgn=1) :
1120	'''Handles NEMO domain changes between NEMO 4.0 to NEMO 4.2
1121
1122	Periodicity and north fold halos has been removed in NEMO 4.2
1123	This routine removes the halos if needed
1124
1125	ptab : Input array (works
1126	rank 2 at least : ptab[...., lat, lon]
1127	nperio : Type of periodicity
1128
1129	See NEMO documentation for further details
1130	'''
1131	nperio = lbc_init (ptab, nperio)[-1]
1132	#lshape = get_shape (ptab)
1133	ax = __find_axis__ (ptab, 'x')[0]
1134	ay = __find_axis__ (ptab, 'y')[0]
1135
1136	if nperio in [4.2, 6.2] :
1137	if ax or ay :
1138	if ax and ay :
1139	ztab = lbc (ptab[..., :-1, 1:-1],
1140	nperio=nperio, cd_type=cd_type, psgn=psgn)
1141	else :
1142	if ax :
1143	ztab = lbc (ptab[..., 1:-1],
1144	nperio=nperio, cd_type=cd_type, psgn=psgn)
1145	if ay :
1146	ztab = lbc (ptab[..., -1],
1147	nperio=nperio, cd_type=cd_type, psgn=psgn)
1148	else :
1149	ztab = ptab
1150	else :
1151	ztab = ptab
1152
1153	return ztab
1154
1155	def lbc_index (jj, ii, jpj, jpi, nperio=None, cd_type='T') :
1156	'''For indexes of a NEMO point, give the corresponding point
1157	inside the domain (i.e. not in the halo)
1158
1159	jj, ii : indexes
1160	jpi, jpi : size of domain
1161	nperio : type of periodicity
1162	cd_type : grid specification : T, U, V or F
1163
1164	See NEMO documentation for further details
1165	'''
1166
1167	if nperio is None :
1168	nperio = __guess_nperio__ (jpj, jpi, nperio)
1169
1170	## For the sake of simplicity, switch to the convention of original
1171	## lbc Fortran routine from NEMO : starts indexes at 1
1172	jy = jj + 1
1173	ix = ii + 1
1174
1175	mmath = __mmath__ (jj)
1176	if mmath is None :
1177	mmath=np
1178
1179	#
1180	#> East-West boundary conditions
1181	# ------------------------------
1182	if nperio in [1, 4, 6] :
1183	#... cyclic
1184	ix = mmath.where (ix==jpi, 2 , ix)
1185	ix = mmath.where (ix== 1 ,jpi-1, ix)
1186
1187	#
1188	def mod_ij (cond, jy_new, ix_new) :
1189	jy_r = mmath.where (cond, jy_new, jy)
1190	ix_r = mmath.where (cond, ix_new, ix)
1191	return jy_r, ix_r
1192	#
1193	#> North-South boundary conditions
1194	# --------------------------------
1195	if nperio in [ 3 , 4 ] :
1196	if cd_type in [ 'T' , 'W' ] :
1197	jy, ix = mod_ij (np.logical_and (jy==jpj , ix>=2 ), jpj-2, jpi-ix+2)
1198	jy, ix = mod_ij (np.logical_and (jy==jpj , ix==1 ), jpj-1, 3 )
1199	jy, ix = mod_ij (np.logical_and (jy==jpj-1, ix>=jpi//2+1),
1200	jy , jpi-ix+2)
1201
1202	if cd_type in [ 'U' ] :
1203	jy, ix = mod_ij (np.logical_and (
1204	jy==jpj ,
1205	np.logical_and (ix>=1, ix <= jpi-1) ),
1206	jy , jpi-ix+1)
1207	jy, ix = mod_ij (np.logical_and (jy==jpj , ix==1 ) , jpj-2, 2 )
1208	jy, ix = mod_ij (np.logical_and (jy==jpj , ix==jpi) , jpj-2, jpi-1 )
1209	jy, ix = mod_ij (np.logical_and (jy==jpj-1,
1210	np.logical_and (ix>=jpi//2, ix<=jpi-1)), jy , jpi-ix+1)
1211
1212	if cd_type in [ 'V' ] :
1213	jy, ix = mod_ij (np.logical_and (jy==jpj-1, ix>=2 ), jpj-2, jpi-ix+2)
1214	jy, ix = mod_ij (np.logical_and (jy==jpj , ix>=2 ), jpj-3, jpi-ix+2)
1215	jy, ix = mod_ij (np.logical_and (jy==jpj , ix==1 ), jpj-3, 3 )
1216
1217	if cd_type in [ 'F' ] :
1218	jy, ix = mod_ij (np.logical_and (jy==jpj-1, ix<=jpi-1), jpj-2, jpi-ix+1)
1219	jy, ix = mod_ij (np.logical_and (jy==jpj , ix<=jpi-1), jpj-3, jpi-ix+1)
1220	jy, ix = mod_ij (np.logical_and (jy==jpj , ix==1 ), jpj-3, 2 )
1221	jy, ix = mod_ij (np.logical_and (jy==jpj , ix==jpi ), jpj-3, jpi-1 )
1222
1223	if nperio in [ 5 , 6 ] :
1224	if cd_type in [ 'T' , 'W' ] : # T-, W-point
1225	jy, ix = mod_ij (jy==jpj, jpj-1, jpi-ix+1)
1226
1227	if cd_type in [ 'U' ] : # U-point
1228	jy, ix = mod_ij (np.logical_and (jy==jpj , ix<=jpi-1 ), jpj-1, jpi-ix )
1229	jy, ix = mod_ij (np.logical_and (jy==jpj , ix==jpi ), jpi-1, 1 )
1230
1231	if cd_type in [ 'V' ] : # V-point
1232	jy, ix = mod_ij (jy==jpj , jy , jpi-ix+1)
1233	jy, ix = mod_ij (np.logical_and (jy==jpj-1, ix>=jpi//2+1), jy , jpi-ix+1)
1234
1235	if cd_type in [ 'F' ] : # F-point
1236	jy, ix = mod_ij (np.logical_and (jy==jpj , ix<=jpi-1 ), jpj-2, jpi-ix )
1237	jy, ix = mod_ij (np.logical_and (ix==jpj , ix==jpi ), jpj-2, 1 )
1238	jy, ix = mod_ij (np.logical_and (jy==jpj-1, ix>=jpi//2+1), jy , jpi-ix )
1239
1240	## Restore convention to Python/C : indexes start at 0
1241	jy += -1
1242	ix += -1
1243
1244	if isinstance (jj, int) :
1245	jy = jy.item ()
1246	if isinstance (ii, int) :
1247	ix = ix.item ()
1248
1249	return jy, ix
1250
1251	def find_ji (lat_data, lon_data, lat_grid, lon_grid, mask=1.0, verbose=False, out=None) :
1252	'''
1253	Description: seeks J,I indices of the grid point which is the closest
1254	of a given point
1255
1256	Usage: go FindJI <data latitude> <data longitude> <grid latitudes> <grid longitudes> [mask]
1257	<grid latitudes><grid longitudes> are 2D fields on J/I (Y/X) dimensions
1258	mask : if given, seek only non masked grid points (i.e with mask=1)
1259
1260	Example : findIJ (40, -20, nav_lat, nav_lon, mask=1.0)
1261
1262	Note : all longitudes and latitudes in degrees
1263
1264	Note : may work with 1D lon/lat (?)
1265	'''
1266	# Get grid dimensions
1267	if len (lon_grid.shape) == 2 :
1268	jpi = lon_grid.shape[-1]
1269	else :
1270	jpi = len(lon_grid)
1271
1272	#mmath = __mmath__ (lat_grid)
1273
1274	# Compute distance from the point to all grid points (in RADian)
1275	arg = ( np.sin (RADlat_data) np.sin (RAD*lat_grid)
1276	+ np.cos (RADlat_data) np.cos (RADlat_grid)
1277	np.cos(RAD*(lon_data-lon_grid)) )
1278	# Send masked points to 'infinite'
1279	distance = np.arccos (arg) + 4.0RPI(1.0-mask)
1280
1281	# Truncates to alleviate some precision problem with some grids
1282	prec = int (1E7)
1283	distance = (distance*prec).astype(int) / prec
1284
1285	# Compute minimum of distance, and index of minimum
1286	#
1287	#distance_min = distance.min ()
1288	jimin = int (distance.argmin ())
1289
1290	# Compute 2D indices (Python/C flavor : starting at 0)
1291	jmin = jimin // jpi
1292	imin = jimin - jmin*jpi
1293
1294	# Result
1295	if verbose :
1296	# Compute distance achieved
1297	#mindist = distance [jmin, imin]
1298
1299	# Compute azimuth
1300	dlon = lon_data-lon_grid[jmin,imin]
1301	arg = np.sin (RAD*dlon) / (
1302	np.cos(RADlat_data)np.tan(RAD*lat_grid[jmin,imin])
1303	- np.sin(RADlat_data)np.cos(RAD*dlon))
1304	azimuth = DAR*np.arctan (arg)
1305	print ( f'I={imin:d} J={jmin:d} - Data:{lat_data:5.1f}Â°N {lon_data:5.1f}Â°E' \
1306	+ f'- Grid:{lat_grid[jmin,imin]:4.1f}Â°N ' \
1307	+ f'{lon_grid[jmin,imin]:4.1f}Â°E - Dist: {RA*distance[jmin,imin]:6.1f}km' \
1308	+ f' {DAR*distance[jmin,imin]:5.2f}Â° ' \
1309	+ f'- Azimuth: {RAD*azimuth:3.2f}RAD - {azimuth:5.1f}Â°' )
1310
1311	if out=='dict' :
1312	return {'x':imin, 'y':jmin}
1313	elif out in ['array', 'numpy', 'np'] :
1314	return np.array ( [jmin, imin] )
1315	elif out in ['xarray', 'xr'] :
1316	return xr.DataArray ( [jmin, imin] )
1317	elif out=='list' :
1318	return [jmin, imin]
1319	elif out=='tuple' :
1320	return jmin, imin
1321	else :
1322	return jmin, imin
1323
1324	def curl (tx, ty, e1f, e2f, nperio=None) :
1325	'''Returns curl of a vector field
1326	'''
1327	ax = __find_axis__ (tx, 'x')[0]
1328	ay = __find_axis__ (ty, 'y')[0]
1329
1330	tx_0 = lbc_add (tx , nperio=nperio, cd_type='U', psgn=-1)
1331	ty_0 = lbc_add (ty , nperio=nperio, cd_type='V', psgn=-1)
1332	e1f_0 = lbc_add (e1f, nperio=nperio, cd_type='U', psgn=-1)
1333	e2f_0 = lbc_add (e2f, nperio=nperio, cd_type='V', psgn=-1)
1334
1335	tx_1 = tx_0.roll ( {ay:-1} )
1336	ty_1 = ty_0.roll ( {ax:-1} )
1337	tx_1 = lbc (tx_1, nperio=nperio, cd_type='U', psgn=-1)
1338	ty_1 = lbc (ty_1, nperio=nperio, cd_type='V', psgn=-1)
1339
1340	zcurl = (ty_1 - ty_0)/e1f_0 - (tx_1 - tx_0)/e2f_0
1341
1342	mask = np.logical_or (
1343	np.logical_or ( np.isnan(tx_0), np.isnan(tx_1)),
1344	np.logical_or ( np.isnan(ty_0), np.isnan(ty_1)) )
1345
1346	zcurl = zcurl.where (np.logical_not (mask), np.nan)
1347
1348	zcurl = lbc_del (zcurl, nperio=nperio, cd_type='F', psgn=1)
1349	zcurl = lbc (zcurl, nperio=nperio, cd_type='F', psgn=1)
1350
1351	return zcurl
1352
1353	def div (ux, uy, e1t, e2t, nperio=None) :
1354	'''Returns divergence of a vector field
1355	'''
1356	ax = __find_axis__ (ux, 'x')[0]
1357	ay = __find_axis__ (ux, 'y')[0]
1358
1359	ux_0 = lbc_add (ux , nperio=nperio, cd_type='U', psgn=-1)
1360	uy_0 = lbc_add (uy , nperio=nperio, cd_type='V', psgn=-1)
1361	e1t_0 = lbc_add (e1t, nperio=nperio, cd_type='U', psgn=-1)
1362	e2t_0 = lbc_add (e2t, nperio=nperio, cd_type='V', psgn=-1)
1363
1364	ux_1 = ux_0.roll ( {ay:1} )
1365	uy_1 = uy_0.roll ( {ax:1} )
1366	ux_1 = lbc (ux_1, nperio=nperio, cd_type='U', psgn=-1)
1367	uy_1 = lbc (uy_1, nperio=nperio, cd_type='V', psgn=-1)
1368
1369	zdiv = (ux_0 - ux_1)/e2t_0 + (uy_0 - uy_1)/e1t_0
1370
1371	mask = np.logical_or (
1372	np.logical_or ( np.isnan(ux_0), np.isnan(ux_1)),
1373	np.logical_or ( np.isnan(uy_0), np.isnan(uy_1)) )
1374
1375	zdiv = zdiv.where (np.logical_not (mask), np.nan)
1376
1377	zdiv = lbc_del (zdiv, nperio=nperio, cd_type='T', psgn=1)
1378	zdiv = lbc (zdiv, nperio=nperio, cd_type='T', psgn=1)
1379
1380	return zdiv
1381
1382	def geo2en (pxx, pyy, pzz, glam, gphi) :
1383	'''Change vector from geocentric to east/north
1384
1385	Inputs :
1386	pxx, pyy, pzz : components on the geocentric system
1387	glam, gphi : longitude and latitude of the points
1388	'''
1389
1390	gsinlon = np.sin (RAD * glam)
1391	gcoslon = np.cos (RAD * glam)
1392	gsinlat = np.sin (RAD * gphi)
1393	gcoslat = np.cos (RAD * gphi)
1394
1395	pte = - pxx * gsinlon + pyy * gcoslon
1396	ptn = - pxx * gcoslon * gsinlat - pyy * gsinlon * gsinlat + pzz * gcoslat
1397
1398	return pte, ptn
1399
1400	def en2geo (pte, ptn, glam, gphi) :
1401	'''Change vector from east/north to geocentric
1402
1403	Inputs :
1404	pte, ptn : eastward/northward components
1405	glam, gphi : longitude and latitude of the points
1406	'''
1407
1408	gsinlon = np.sin (RAD * glam)
1409	gcoslon = np.cos (RAD * glam)
1410	gsinlat = np.sin (RAD * gphi)
1411	gcoslat = np.cos (RAD * gphi)
1412
1413	pxx = - pte * gsinlon - ptn * gcoslon * gsinlat
1414	pyy = pte * gcoslon - ptn * gsinlon * gsinlat
1415	pzz = ptn * gcoslat
1416
1417	return pxx, pyy, pzz
1418
1419
1420	def clo_lon (lon, lon0=0., rad=False, deg=True) :
1421	'''Choose closest to lon0 longitude, adding/substacting 360Â°
1422	if needed
1423	'''
1424	mmath = __mmath__ (lon, np)
1425	if rad :
1426	lon_range = 2.*np.pi
1427	if deg :
1428	lon_range = 360.
1429	c_lon = lon
1430	c_lon = mmath.where (c_lon > lon0 + lon_range*0.5,
1431	c_lon-lon_range, clo_lon)
1432	c_lon = mmath.where (c_lon < lon0 - lon_range*0.5,
1433	c_lon+lon_range, clo_lon)
1434	c_lon = mmath.where (c_lon > lon0 + lon_range*0.5,
1435	c_lon-lon_range, clo_lon)
1436	c_lon = mmath.where (c_lon < lon0 - lon_range*0.5,
1437	c_lon+lon_range, clo_lon)
1438	if c_lon.shape == () :
1439	c_lon = c_lon.item ()
1440	if mmath == xr :
1441	if lon.attrs :
1442	c_lon.attrs.update ( lon.attrs )
1443	return c_lon
1444
1445	def index2depth (pk, gdept_0) :
1446	'''From index (real, continuous), get depth
1447
1448	Needed to use transforms in Matplotlib
1449	'''
1450	jpk = gdept_0.shape[0]
1451	kk = xr.DataArray(pk)
1452	k = np.maximum (0, np.minimum (jpk-1, kk ))
1453	k0 = np.floor (k).astype (int)
1454	k1 = np.maximum (0, np.minimum (jpk-1, k0+1))
1455	zz = k - k0
1456	gz = (1.0-zz)gdept_0[k0]+ zzgdept_0[k1]
1457	return gz.values
1458
1459	def depth2index (pz, gdept_0) :
1460	'''From depth, get index (real, continuous)
1461
1462	Needed to use transforms in Matplotlib
1463	'''
1464	jpk = gdept_0.shape[0]
1465	if isinstance (pz, xr.core.dataarray.DataArray ) :
1466	zz = xr.DataArray (pz.values, dims=('zz',))
1467	elif isinstance (pz, np.ndarray) :
1468	zz = xr.DataArray (pz.ravel(), dims=('zz',))
1469	else :
1470	zz = xr.DataArray (np.array([pz]).ravel(), dims=('zz',))
1471	zz = np.minimum (gdept_0[-1], np.maximum (0, zz))
1472
1473	idk1 = np.minimum ( (gdept_0-zz), 0.).argmax (axis=0).astype(int)
1474	idk1 = np.maximum (0, np.minimum (jpk-1, idk1 ))
1475	idk2 = np.maximum (0, np.minimum (jpk-1, idk1-1))
1476
1477	zff = (zz - gdept_0[idk2])/(gdept_0[idk1]-gdept_0[idk2])
1478	idk = zffidk1 + (1.0-zff)idk2
1479	idk = xr.where ( np.isnan(idk), idk1, idk)
1480	return idk.values
1481
1482	def index2depth_panels (pk, gdept_0, depth0, fact) :
1483	'''From index (real, continuous), get depth, with bottom part compressed
1484
1485	Needed to use transforms in Matplotlib
1486	'''
1487	jpk = gdept_0.shape[0]
1488	kk = xr.DataArray (pk)
1489	k = np.maximum (0, np.minimum (jpk-1, kk ))
1490	k0 = np.floor (k).astype (int)
1491	k1 = np.maximum (0, np.minimum (jpk-1, k0+1))
1492	zz = k - k0
1493	gz = (1.0-zz)gdept_0[k0]+ zzgdept_0[k1]
1494	gz = xr.where ( gz<depth0, gz, depth0 + (gz-depth0)*fact)
1495	return gz.values
1496
1497	def depth2index_panels (pz, gdept_0, depth0, fact) :
1498	'''From index (real, continuous), get depth, with bottom part compressed
1499
1500	Needed to use transforms in Matplotlib
1501	'''
1502	jpk = gdept_0.shape[0]
1503	if isinstance (pz, xr.core.dataarray.DataArray) :
1504	zz = xr.DataArray (pz.values , dims=('zz',))
1505	elif isinstance (pz, np.ndarray) :
1506	zz = xr.DataArray (pz.ravel(), dims=('zz',))
1507	else :
1508	zz = xr.DataArray (np.array([pz]).ravel(), dims=('zz',))
1509	zz = np.minimum (gdept_0[-1], np.maximum (0, zz))
1510	gdept_comp = xr.where ( gdept_0>depth0,
1511	(gdept_0-depth0)*fact+depth0, gdept_0)
1512	zz_comp = xr.where ( zz >depth0, (zz -depth0)*fact+depth0,
1513	zz )
1514	zz_comp = np.minimum (gdept_comp[-1], np.maximum (0, zz_comp))
1515
1516	idk1 = np.minimum ( (gdept_0-zz_comp), 0.).argmax (axis=0).astype(int)
1517	idk1 = np.maximum (0, np.minimum (jpk-1, idk1 ))
1518	idk2 = np.maximum (0, np.minimum (jpk-1, idk1-1))
1519
1520	zff = (zz_comp - gdept_0[idk2])/(gdept_0[idk1]-gdept_0[idk2])
1521	idk = zffidk1 + (1.0-zff)idk2
1522	idk = xr.where ( np.isnan(idk), idk1, idk)
1523	return idk.values
1524
1525	def depth2comp (pz, depth0, fact ) :
1526	'''Form depth, get compressed depth, with bottom part compressed
1527
1528	Needed to use transforms in Matplotlib
1529	'''
1530	#print ('start depth2comp')
1531	if isinstance (pz, xr.core.dataarray.DataArray) :
1532	zz = pz.values
1533	elif isinstance(pz, list) :
1534	zz = np.array (pz)
1535	else :
1536	zz = pz
1537	gz = np.where ( zz>depth0, (zz-depth0)*fact+depth0, zz)
1538	#print ( f'depth2comp : {gz=}' )
1539	if type (pz) in [int, float] :
1540	return gz.item()
1541	else :
1542	return gz
1543
1544	def comp2depth (pz, depth0, fact ) :
1545	'''Form compressed depth, get depth, with bottom part compressed
1546
1547	Needed to use transforms in Matplotlib
1548	'''
1549	if isinstance (pz, xr.core.dataarray.DataArray) :
1550	zz = pz.values
1551	elif isinstance (pz, list) :
1552	zz = np.array (pz)
1553	else :
1554	zz = pz
1555	gz = np.where ( zz>depth0, (zz-depth0)/fact+depth0, zz)
1556	if type (pz) in [int, float] :
1557	gz = gz.item()
1558
1559	return gz
1560
1561	def index2lon (pi, plon_1d) :
1562	'''From index (real, continuous), get longitude
1563
1564	Needed to use transforms in Matplotlib
1565	'''
1566	jpi = plon_1d.shape[0]
1567	ii = xr.DataArray (pi)
1568	i = np.maximum (0, np.minimum (jpi-1, ii ))
1569	i0 = np.floor (i).astype (int)
1570	i1 = np.maximum (0, np.minimum (jpi-1, i0+1))
1571	xx = i - i0
1572	gx = (1.0-xx)plon_1d[i0]+ xxplon_1d[i1]
1573	return gx.values
1574
1575	def lon2index (px, plon_1d) :
1576	'''From longitude, get index (real, continuous)
1577
1578	Needed to use transforms in Matplotlib
1579	'''
1580	jpi = plon_1d.shape[0]
1581	if isinstance (px, xr.core.dataarray.DataArray) :
1582	xx = xr.DataArray (px.values , dims=('xx',))
1583	elif isinstance (px, np.ndarray) :
1584	xx = xr.DataArray (px.ravel(), dims=('xx',))
1585	else :
1586	xx = xr.DataArray (np.array([px]).ravel(), dims=('xx',))
1587	xx = xr.where ( xx>plon_1d.max(), xx-360.0, xx)
1588	xx = xr.where ( xx<plon_1d.min(), xx+360.0, xx)
1589	xx = np.minimum (plon_1d.max(), np.maximum(xx, plon_1d.min() ))
1590	idi1 = np.minimum ( (plon_1d-xx), 0.).argmax (axis=0).astype(int)
1591	idi1 = np.maximum (0, np.minimum (jpi-1, idi1 ))
1592	idi2 = np.maximum (0, np.minimum (jpi-1, idi1-1))
1593
1594	zff = (xx - plon_1d[idi2])/(plon_1d[idi1]-plon_1d[idi2])
1595	idi = zffidi1 + (1.0-zff)idi2
1596	idi = xr.where ( np.isnan(idi), idi1, idi)
1597	return idi.values
1598
1599	def index2lat (pj, plat_1d) :
1600	'''From index (real, continuous), get latitude
1601
1602	Needed to use transforms in Matplotlib
1603	'''
1604	jpj = plat_1d.shape[0]
1605	jj = xr.DataArray (pj)
1606	j = np.maximum (0, np.minimum (jpj-1, jj ))
1607	j0 = np.floor (j).astype (int)
1608	j1 = np.maximum (0, np.minimum (jpj-1, j0+1))
1609	yy = j - j0
1610	gy = (1.0-yy)plat_1d[j0]+ yyplat_1d[j1]
1611	return gy.values
1612
1613	def lat2index (py, plat_1d) :
1614	'''From latitude, get index (real, continuous)
1615
1616	Needed to use transforms in Matplotlib
1617	'''
1618	jpj = plat_1d.shape[0]
1619	if isinstance (py, xr.core.dataarray.DataArray) :
1620	yy = xr.DataArray (py.values , dims=('yy',))
1621	elif isinstance (py, np.ndarray) :
1622	yy = xr.DataArray (py.ravel(), dims=('yy',))
1623	else :
1624	yy = xr.DataArray (np.array([py]).ravel(), dims=('yy',))
1625	yy = np.minimum (plat_1d.max(), np.minimum(yy, plat_1d.max() ))
1626	idj1 = np.minimum ( (plat_1d-yy), 0.).argmax (axis=0).astype(int)
1627	idj1 = np.maximum (0, np.minimum (jpj-1, idj1 ))
1628	idj2 = np.maximum (0, np.minimum (jpj-1, idj1-1))
1629
1630	zff = (yy - plat_1d[idj2])/(plat_1d[idj1]-plat_1d[idj2])
1631	idj = zffidj1 + (1.0-zff)idj2
1632	idj = xr.where ( np.isnan(idj), idj1, idj)
1633	return idj.values
1634
1635	def angle_full (glamt, gphit, glamu, gphiu, glamv, gphiv,
1636	glamf, gphif, nperio=None) :
1637	'''Computes sinus and cosinus of model line direction with
1638	respect to east
1639	'''
1640	mmath = __mmath__ (glamt)
1641
1642	zlamt = lbc_add (glamt, nperio, 'T', 1.)
1643	zphit = lbc_add (gphit, nperio, 'T', 1.)
1644	zlamu = lbc_add (glamu, nperio, 'U', 1.)
1645	zphiu = lbc_add (gphiu, nperio, 'U', 1.)
1646	zlamv = lbc_add (glamv, nperio, 'V', 1.)
1647	zphiv = lbc_add (gphiv, nperio, 'V', 1.)
1648	zlamf = lbc_add (glamf, nperio, 'F', 1.)
1649	zphif = lbc_add (gphif, nperio, 'F', 1.)
1650
1651	# north pole direction & modulous (at T-point)
1652	zxnpt = 0. - 2.0 * np.cos (RADzlamt) np.tan (RPI/4.0 - RAD*zphit/2.0)
1653	zynpt = 0. - 2.0 * np.sin (RADzlamt) np.tan (RPI/4.0 - RAD*zphit/2.0)
1654	znnpt = zxnptzxnpt + zynptzynpt
1655
1656	# north pole direction & modulous (at U-point)
1657	zxnpu = 0. - 2.0 * np.cos (RADzlamu) np.tan (RPI/4.0 - RAD*zphiu/2.0)
1658	zynpu = 0. - 2.0 * np.sin (RADzlamu) np.tan (RPI/4.0 - RAD*zphiu/2.0)
1659	znnpu = zxnpuzxnpu + zynpuzynpu
1660
1661	# north pole direction & modulous (at V-point)
1662	zxnpv = 0. - 2.0 * np.cos (RADzlamv) np.tan (RPI/4.0 - RAD*zphiv/2.0)
1663	zynpv = 0. - 2.0 * np.sin (RADzlamv) np.tan (RPI/4.0 - RAD*zphiv/2.0)
1664	znnpv = zxnpvzxnpv + zynpvzynpv
1665
1666	# north pole direction & modulous (at F-point)
1667	zxnpf = 0. - 2.0 * np.cos( RADzlamf ) np.tan ( RPI/4. - RAD*zphif/2. )
1668	zynpf = 0. - 2.0 * np.sin( RADzlamf ) np.tan ( RPI/4. - RAD*zphif/2. )
1669	znnpf = zxnpfzxnpf + zynpfzynpf
1670
1671	# j-direction: v-point segment direction (around T-point)
1672	zlam = zlamv
1673	zphi = zphiv
1674	zlan = np.roll ( zlamv, axis=-2, shift=1) # glamv (ji,jj-1)
1675	zphh = np.roll ( zphiv, axis=-2, shift=1) # gphiv (ji,jj-1)
1676	zxvvt = 2.0 * np.cos ( RADzlam ) np.tan ( RPI/4. - RAD*zphi/2. ) \
1677	- 2.0 * np.cos ( RADzlan ) np.tan ( RPI/4. - RAD*zphh/2. )
1678	zyvvt = 2.0 * np.sin ( RADzlam ) np.tan ( RPI/4. - RAD*zphi/2. ) \
1679	- 2.0 * np.sin ( RADzlan ) np.tan ( RPI/4. - RAD*zphh/2. )
1680	znvvt = np.sqrt ( znnpt * ( zxvvtzxvvt + zyvvtzyvvt ) )
1681
1682	# j-direction: f-point segment direction (around u-point)
1683	zlam = zlamf
1684	zphi = zphif
1685	zlan = np.roll (zlamf, axis=-2, shift=1) # glamf (ji,jj-1)
1686	zphh = np.roll (zphif, axis=-2, shift=1) # gphif (ji,jj-1)
1687	zxffu = 2.0 * np.cos ( RADzlam ) np.tan ( RPI/4. - RAD*zphi/2. ) \
1688	- 2.0 * np.cos ( RADzlan ) np.tan ( RPI/4. - RAD*zphh/2. )
1689	zyffu = 2.0 * np.sin ( RADzlam ) np.tan ( RPI/4. - RAD*zphi/2. ) \
1690	- 2.0 * np.sin ( RADzlan ) np.tan ( RPI/4. - RAD*zphh/2. )
1691	znffu = np.sqrt ( znnpu * ( zxffuzxffu + zyffuzyffu ) )
1692
1693	# i-direction: f-point segment direction (around v-point)
1694	zlam = zlamf
1695	zphi = zphif
1696	zlan = np.roll (zlamf, axis=-1, shift=1) # glamf (ji-1,jj)
1697	zphh = np.roll (zphif, axis=-1, shift=1) # gphif (ji-1,jj)
1698	zxffv = 2.0 * np.cos ( RADzlam ) np.tan ( RPI/4. - RAD*zphi/2. ) \
1699	- 2.0 * np.cos ( RADzlan ) np.tan ( RPI/4. - RAD*zphh/2. )
1700	zyffv = 2.0 * np.sin ( RADzlam ) np.tan ( RPI/4. - RAD*zphi/2. ) \
1701	- 2.0 * np.sin ( RADzlan ) np.tan ( RPI/4. - RAD*zphh/2. )
1702	znffv = np.sqrt ( znnpv * ( zxffvzxffv + zyffvzyffv ) )
1703
1704	# j-direction: u-point segment direction (around f-point)
1705	zlam = np.roll (zlamu, axis=-2, shift=-1) # glamu (ji,jj+1)
1706	zphi = np.roll (zphiu, axis=-2, shift=-1) # gphiu (ji,jj+1)
1707	zlan = zlamu
1708	zphh = zphiu
1709	zxuuf = 2. * np.cos ( RADzlam ) np.tan ( RPI/4. - RAD*zphi/2. ) \
1710	- 2. * np.cos ( RADzlan ) np.tan ( RPI/4. - RAD*zphh/2. )
1711	zyuuf = 2. * np.sin ( RADzlam ) np.tan ( RPI/4. - RAD*zphi/2. ) \
1712	- 2. * np.sin ( RADzlan ) np.tan ( RPI/4. - RAD*zphh/2. )
1713	znuuf = np.sqrt ( znnpf * ( zxuufzxuuf + zyuufzyuuf ) )
1714
1715
1716	# cosinus and sinus using scalar and vectorial products
1717	gsint = ( zxnptzyvvt - zynptzxvvt ) / znvvt
1718	gcost = ( zxnptzxvvt + zynptzyvvt ) / znvvt
1719
1720	gsinu = ( zxnpuzyffu - zynpuzxffu ) / znffu
1721	gcosu = ( zxnpuzxffu + zynpuzyffu ) / znffu
1722
1723	gsinf = ( zxnpfzyuuf - zynpfzxuuf ) / znuuf
1724	gcosf = ( zxnpfzxuuf + zynpfzyuuf ) / znuuf
1725
1726	gsinv = ( zxnpvzxffv + zynpvzyffv ) / znffv
1727	# (caution, rotation of 90 degres)
1728	gcosv =-( zxnpvzyffv - zynpvzxffv ) / znffv
1729
1730	gsint = lbc_del (gsint, cd_type='T', nperio=nperio, psgn=-1.)
1731	gcost = lbc_del (gcost, cd_type='T', nperio=nperio, psgn=-1.)
1732	gsinu = lbc_del (gsinu, cd_type='U', nperio=nperio, psgn=-1.)
1733	gcosu = lbc_del (gcosu, cd_type='U', nperio=nperio, psgn=-1.)
1734	gsinv = lbc_del (gsinv, cd_type='V', nperio=nperio, psgn=-1.)
1735	gcosv = lbc_del (gcosv, cd_type='V', nperio=nperio, psgn=-1.)
1736	gsinf = lbc_del (gsinf, cd_type='F', nperio=nperio, psgn=-1.)
1737	gcosf = lbc_del (gcosf, cd_type='F', nperio=nperio, psgn=-1.)
1738
1739	if mmath == xr :
1740	gsint = gsint.assign_coords ( glamt.coords )
1741	gcost = gcost.assign_coords ( glamt.coords )
1742	gsinu = gsinu.assign_coords ( glamu.coords )
1743	gcosu = gcosu.assign_coords ( glamu.coords )
1744	gsinv = gsinv.assign_coords ( glamv.coords )
1745	gcosv = gcosv.assign_coords ( glamv.coords )
1746	gsinf = gsinf.assign_coords ( glamf.coords )
1747	gcosf = gcosf.assign_coords ( glamf.coords )
1748
1749	return gsint, gcost, gsinu, gcosu, gsinv, gcosv, gsinf, gcosf
1750
1751	def angle (glam, gphi, nperio, cd_type='T') :
1752	'''Computes sinus and cosinus of model line direction with
1753	respect to east
1754	'''
1755	mmath = __mmath__ (glam)
1756
1757	zlam = lbc_add (glam, nperio, cd_type, 1.)
1758	zphi = lbc_add (gphi, nperio, cd_type, 1.)
1759
1760	# north pole direction & modulous
1761	zxnp = 0. - 2.0 * np.cos (RADzlam) np.tan (RPI/4.0 - RAD*zphi/2.0)
1762	zynp = 0. - 2.0 * np.sin (RADzlam) np.tan (RPI/4.0 - RAD*zphi/2.0)
1763	znnp = zxnpzxnp + zynpzynp
1764
1765	# j-direction: segment direction (around point)
1766	zlan_n = np.roll (zlam, axis=-2, shift=-1) # glam [jj+1, ji]
1767	zphh_n = np.roll (zphi, axis=-2, shift=-1) # gphi [jj+1, ji]
1768	zlan_s = np.roll (zlam, axis=-2, shift= 1) # glam [jj-1, ji]
1769	zphh_s = np.roll (zphi, axis=-2, shift= 1) # gphi [jj-1, ji]
1770
1771	zxff = 2.0 * np.cos (RADzlan_n) np.tan (RPI/4.0 - RAD*zphh_n/2.0) \
1772	- 2.0 * np.cos (RADzlan_s) np.tan (RPI/4.0 - RAD*zphh_s/2.0)
1773	zyff = 2.0 * np.sin (RADzlan_n) np.tan (RPI/4.0 - RAD*zphh_n/2.0) \
1774	- 2.0 * np.sin (RADzlan_s) np.tan (RPI/4.0 - RAD*zphh_s/2.0)
1775	znff = np.sqrt (znnp * (zxffzxff + zyffzyff) )
1776
1777	gsin = (zxnpzyff - zynpzxff) / znff
1778	gcos = (zxnpzxff + zynpzyff) / znff
1779
1780	gsin = lbc_del (gsin, cd_type=cd_type, nperio=nperio, psgn=-1.)
1781	gcos = lbc_del (gcos, cd_type=cd_type, nperio=nperio, psgn=-1.)
1782
1783	if mmath == xr :
1784	gsin = gsin.assign_coords ( glam.coords )
1785	gcos = gcos.assign_coords ( glam.coords )
1786
1787	return gsin, gcos
1788
1789	def rot_en2ij ( u_e, v_n, gsin, gcos, nperio, cd_type ) :
1790	'''Rotates the Repere: Change vector componantes between
1791	geographic grid --> stretched coordinates grid.
1792
1793	All components are on the same grid (T, U, V or F)
1794	'''
1795
1796	u_i = + u_e * gcos + v_n * gsin
1797	v_j = - u_e * gsin + v_n * gcos
1798
1799	u_i = lbc (u_i, nperio=nperio, cd_type=cd_type, psgn=-1.0)
1800	v_j = lbc (v_j, nperio=nperio, cd_type=cd_type, psgn=-1.0)
1801
1802	return u_i, v_j
1803
1804	def rot_ij2en ( u_i, v_j, gsin, gcos, nperio, cd_type='T' ) :
1805	'''Rotates the Repere: Change vector componantes from
1806	stretched coordinates grid --> geographic grid
1807
1808	All components are on the same grid (T, U, V or F)
1809	'''
1810	u_e = + u_i * gcos - v_j * gsin
1811	v_n = + u_i * gsin + v_j * gcos
1812
1813	u_e = lbc (u_e, nperio=nperio, cd_type=cd_type, psgn=1.0)
1814	v_n = lbc (v_n, nperio=nperio, cd_type=cd_type, psgn=1.0)
1815
1816	return u_e, v_n
1817
1818	def rot_uv2en ( uo, vo, gsint, gcost, nperio, zdim=None ) :
1819	'''Rotate the Repere: Change vector componantes from
1820	stretched coordinates grid --> geographic grid
1821
1822	uo : velocity along i at the U grid point
1823	vo : valocity along j at the V grid point
1824
1825	Returns east-north components on the T grid point
1826	'''
1827	ut = u2t (uo, nperio=nperio, psgn=-1.0, zdim=zdim)
1828	vt = v2t (vo, nperio=nperio, psgn=-1.0, zdim=zdim)
1829
1830	u_e = + ut * gcost - vt * gsint
1831	v_n = + ut * gsint + vt * gcost
1832
1833	u_e = lbc (u_e, nperio=nperio, cd_type='T', psgn=1.0)
1834	v_n = lbc (v_n, nperio=nperio, cd_type='T', psgn=1.0)
1835
1836	return u_e, v_n
1837
1838	def rot_uv2enf ( uo, vo, gsinf, gcosf, nperio, zdim=None ) :
1839	'''Rotates the Repere: Change vector componantes from
1840	stretched coordinates grid --> geographic grid
1841
1842	uo : velocity along i at the U grid point
1843	vo : valocity along j at the V grid point
1844
1845	Returns east-north components on the F grid point
1846	'''
1847	uf = u2f (uo, nperio=nperio, psgn=-1.0, zdim=zdim)
1848	vf = v2f (vo, nperio=nperio, psgn=-1.0, zdim=zdim)
1849
1850	u_e = + uf * gcosf - vf * gsinf
1851	v_n = + uf * gsinf + vf * gcosf
1852
1853	u_e = lbc (u_e, nperio=nperio, cd_type='F', psgn= 1.0)
1854	v_n = lbc (v_n, nperio=nperio, cd_type='F', psgn= 1.0)
1855
1856	return u_e, v_n
1857
1858	def u2t (utab, nperio=None, psgn=-1.0, zdim=None, action='ave') :
1859	'''Interpolates an array from U grid to T grid (i-mean)
1860	'''
1861	mmath = __mmath__ (utab)
1862	utab_0 = mmath.where ( np.isnan(utab), 0., utab)
1863	#lperio, aperio = lbc_diag (nperio)
1864	utab_0 = lbc_add (utab_0, nperio=nperio, cd_type='U', psgn=psgn)
1865	ax, ix = __find_axis__ (utab_0, 'x')
1866	az = __find_axis__ (utab_0, 'z')[0]
1867
1868	if ax :
1869	if action == 'ave' :
1870	ttab = 0.5 * (utab_0 + np.roll (utab_0, axis=ix, shift=1))
1871	if action == 'min' :
1872	ttab = np.minimum (utab_0 , np.roll (utab_0, axis=ix, shift=1))
1873	if action == 'max' :
1874	ttab = np.maximum (utab_0 , np.roll (utab_0, axis=ix, shift=1))
1875	if action == 'mult':
1876	ttab = utab_0 * np.roll (utab_0, axis=ix, shift=1)
1877	ttab = lbc_del (ttab , nperio=nperio, cd_type='T', psgn=psgn)
1878	else :
1879	ttab = lbc_del (utab_0, nperio=nperio, cd_type='T', psgn=psgn)
1880
1881	if mmath == xr :
1882	if ax :
1883	ttab = ttab.assign_coords({ax:np.arange (ttab.shape[ix])+1.})
1884	if zdim and az :
1885	if az != zdim :
1886	ttab = ttab.rename( {az:zdim})
1887	return ttab
1888
1889	def v2t (vtab, nperio=None, psgn=-1.0, zdim=None, action='ave') :
1890	'''Interpolates an array from V grid to T grid (j-mean)
1891	'''
1892	mmath = __mmath__ (vtab)
1893	#lperio, aperio = lbc_diag (nperio)
1894	vtab_0 = mmath.where ( np.isnan(vtab), 0., vtab)
1895	vtab_0 = lbc_add (vtab_0, nperio=nperio, cd_type='V', psgn=psgn)
1896	ay, jy = __find_axis__ (vtab_0, 'y')
1897	az = __find_axis__ (vtab_0, 'z')[0]
1898	if ay :
1899	if action == 'ave' :
1900	ttab = 0.5 * (vtab_0 + np.roll (vtab_0, axis=jy, shift=1))
1901	if action == 'min' :
1902	ttab = np.minimum (vtab_0 , np.roll (vtab_0, axis=jy, shift=1))
1903	if action == 'max' :
1904	ttab = np.maximum (vtab_0 , np.roll (vtab_0, axis=jy, shift=1))
1905	if action == 'mult' :
1906	ttab = vtab_0 * np.roll (vtab_0, axis=jy, shift=1)
1907	ttab = lbc_del (ttab , nperio=nperio, cd_type='T', psgn=psgn)
1908	else :
1909	ttab = lbc_del (vtab_0, nperio=nperio, cd_type='T', psgn=psgn)
1910
1911	if mmath == xr :
1912	if ay :
1913	ttab = ttab.assign_coords({ay:np.arange(ttab.shape[jy])+1.})
1914	if zdim and az :
1915	if az != zdim :
1916	ttab = ttab.rename( {az:zdim})
1917	return ttab
1918
1919	def f2t (ftab, nperio=None, psgn=1.0, zdim=None, action='ave') :
1920	'''Interpolates an array from F grid to T grid (i- and j- means)
1921	'''
1922	mmath = __mmath__ (ftab)
1923	ftab_0 = mmath.where ( np.isnan(ftab), 0., ftab)
1924	ftab_0 = lbc_add (ftab_0 , nperio=nperio, cd_type='F', psgn=psgn)
1925	ttab = v2t (f2v (ftab_0, nperio=nperio, psgn=psgn, zdim=zdim, action=action),
1926	nperio=nperio, psgn=psgn, zdim=zdim, action=action)
1927	return lbc_del (ttab, nperio=nperio, cd_type='T', psgn=psgn)
1928
1929	def t2u (ttab, nperio=None, psgn=1.0, zdim=None, action='ave') :
1930	'''Interpolates an array from T grid to U grid (i-mean)
1931	'''
1932	mmath = __mmath__ (ttab)
1933	ttab_0 = mmath.where ( np.isnan(ttab), 0., ttab)
1934	ttab_0 = lbc_add (ttab_0 , nperio=nperio, cd_type='T', psgn=psgn)
1935	ax, ix = __find_axis__ (ttab_0, 'x')[0]
1936	az = __find_axis__ (ttab_0, 'z')
1937	if ix :
1938	if action == 'ave' :
1939	utab = 0.5 * (ttab_0 + np.roll (ttab_0, axis=ix, shift=-1))
1940	if action == 'min' :
1941	utab = np.minimum (ttab_0 , np.roll (ttab_0, axis=ix, shift=-1))
1942	if action == 'max' :
1943	utab = np.maximum (ttab_0 , np.roll (ttab_0, axis=ix, shift=-1))
1944	if action == 'mult' :
1945	utab = ttab_0 * np.roll (ttab_0, axis=ix, shift=-1)
1946	utab = lbc_del (utab , nperio=nperio, cd_type='U', psgn=psgn)
1947	else :
1948	utab = lbc_del (ttab_0, nperio=nperio, cd_type='U', psgn=psgn)
1949
1950	if mmath == xr :
1951	if ax :
1952	utab = ttab.assign_coords({ax:np.arange(utab.shape[ix])+1.})
1953	if zdim and az :
1954	if az != zdim :
1955	utab = utab.rename( {az:zdim})
1956	return utab
1957
1958	def t2v (ttab, nperio=None, psgn=1.0, zdim=None, action='ave') :
1959	'''Interpolates an array from T grid to V grid (j-mean)
1960	'''
1961	mmath = __mmath__ (ttab)
1962	ttab_0 = mmath.where ( np.isnan(ttab), 0., ttab)
1963	ttab_0 = lbc_add (ttab_0 , nperio=nperio, cd_type='T', psgn=psgn)
1964	ay, jy = __find_axis__ (ttab_0, 'y')
1965	az = __find_axis__ (ttab_0, 'z')[0]
1966	if jy :
1967	if action == 'ave' :
1968	vtab = 0.5 * (ttab_0 + np.roll (ttab_0, axis=jy, shift=-1))
1969	if action == 'min' :
1970	vtab = np.minimum (ttab_0 , np.roll (ttab_0, axis=jy, shift=-1))
1971	if action == 'max' :
1972	vtab = np.maximum (ttab_0 , np.roll (ttab_0, axis=jy, shift=-1))
1973	if action == 'mult' :
1974	vtab = ttab_0 * np.roll (ttab_0, axis=jy, shift=-1)
1975	vtab = lbc_del (vtab , nperio=nperio, cd_type='V', psgn=psgn)
1976	else :
1977	vtab = lbc_del (ttab_0, nperio=nperio, cd_type='V', psgn=psgn)
1978
1979	if mmath == xr :
1980	if ay :
1981	vtab = vtab.assign_coords({ay:np.arange(vtab.shape[jy])+1.})
1982	if zdim and az :
1983	if az != zdim :
1984	vtab = vtab.rename( {az:zdim})
1985	return vtab
1986
1987	def v2f (vtab, nperio=None, psgn=-1.0, zdim=None, action='ave') :
1988	'''Interpolates an array from V grid to F grid (i-mean)
1989	'''
1990	mmath = __mmath__ (vtab)
1991	vtab_0 = mmath.where ( np.isnan(vtab), 0., vtab)
1992	vtab_0 = lbc_add (vtab_0 , nperio=nperio, cd_type='V', psgn=psgn)
1993	ax, ix = __find_axis__ (vtab_0, 'x')
1994	az = __find_axis__ (vtab_0, 'z')[0]
1995	if ix :
1996	if action == 'ave' :
1997	ftab = 0.5 * (vtab_0 + np.roll (vtab_0, axis=ix, shift=-1))
1998	if action == 'min' :
1999	ftab = np.minimum (vtab_0 , np.roll (vtab_0, axis=ix, shift=-1))
2000	if action == 'max' :
2001	ftab = np.maximum (vtab_0 , np.roll (vtab_0, axis=ix, shift=-1))
2002	if action == 'mult' :
2003	ftab = vtab_0 * np.roll (vtab_0, axis=ix, shift=-1)
2004	ftab = lbc_del (ftab , nperio=nperio, cd_type='F', psgn=psgn)
2005	else :
2006	ftab = lbc_del (vtab_0, nperio=nperio, cd_type='F', psgn=psgn)
2007
2008	if mmath == xr :
2009	if ax :
2010	ftab = ftab.assign_coords({ax:np.arange(ftab.shape[ix])+1.})
2011	if zdim and az :
2012	if az != zdim :
2013	ftab = ftab.rename( {az:zdim})
2014	return lbc_del (ftab, nperio=nperio, cd_type='F', psgn=psgn)
2015
2016	def u2f (utab, nperio=None, psgn=-1.0, zdim=None, action='ave') :
2017	'''Interpolates an array from U grid to F grid i-mean)
2018	'''
2019	mmath = __mmath__ (utab)
2020	utab_0 = mmath.where ( np.isnan(utab), 0., utab)
2021	utab_0 = lbc_add (utab_0 , nperio=nperio, cd_type='U', psgn=psgn)
2022	ay, jy = __find_axis__ (utab_0, 'y')
2023	az = __find_axis__ (utab_0, 'z')[0]
2024	if jy :
2025	if action == 'ave' :
2026	ftab = 0.5 * (utab_0 + np.roll (utab_0, axis=jy, shift=-1))
2027	if action == 'min' :
2028	ftab = np.minimum (utab_0 , np.roll (utab_0, axis=jy, shift=-1))
2029	if action == 'max' :
2030	ftab = np.maximum (utab_0 , np.roll (utab_0, axis=jy, shift=-1))
2031	if action == 'mult' :
2032	ftab = utab_0 * np.roll (utab_0, axis=jy, shift=-1)
2033	ftab = lbc_del (ftab, nperio=nperio, cd_type='F', psgn=psgn)
2034	else :
2035	ftab = lbc_del (utab_0, nperio=nperio, cd_type='F', psgn=psgn)
2036
2037	if mmath == xr :
2038	if ay :
2039	ftab = ftab.assign_coords({'y':np.arange(ftab.shape[jy])+1.})
2040	if zdim and az :
2041	if az != zdim :
2042	ftab = ftab.rename( {az:zdim})
2043	return ftab
2044
2045	def t2f (ttab, nperio=None, psgn=1.0, zdim=None, action='mean') :
2046	'''Interpolates an array on T grid to F grid (i- and j- means)
2047	'''
2048	mmath = __mmath__ (ttab)
2049	ttab_0 = mmath.where ( np.isnan(ttab), 0., ttab)
2050	ttab_0 = lbc_add (ttab_0 , nperio=nperio, cd_type='T', psgn=psgn)
2051	ftab = t2u (u2f (ttab, nperio=nperio, psgn=psgn, zdim=zdim, action=action),
2052	nperio=nperio, psgn=psgn, zdim=zdim, action=action)
2053
2054	return lbc_del (ftab, nperio=nperio, cd_type='F', psgn=psgn)
2055
2056	def f2u (ftab, nperio=None, psgn=1.0, zdim=None, action='ave') :
2057	'''Interpolates an array on F grid to U grid (j-mean)
2058	'''
2059	mmath = __mmath__ (ftab)
2060	ftab_0 = mmath.where ( np.isnan(ftab), 0., ftab)
2061	ftab_0 = lbc_add (ftab_0 , nperio=nperio, cd_type='F', psgn=psgn)
2062	ay, jy = __find_axis__ (ftab_0, 'y')
2063	az = __find_axis__ (ftab_0, 'z')[0]
2064	if jy :
2065	if action == 'ave' :
2066	utab = 0.5 * (ftab_0 + np.roll (ftab_0, axis=jy, shift=-1))
2067	if action == 'min' :
2068	utab = np.minimum (ftab_0 , np.roll (ftab_0, axis=jy, shift=-1))
2069	if action == 'max' :
2070	utab = np.maximum (ftab_0 , np.roll (ftab_0, axis=jy, shift=-1))
2071	if action == 'mult' :
2072	utab = ftab_0 * np.roll (ftab_0, axis=jy, shift=-1)
2073	utab = lbc_del (utab , nperio=nperio, cd_type='U', psgn=psgn)
2074	else :
2075	utab = lbc_del (ftab_0, nperio=nperio, cd_type='U', psgn=psgn)
2076
2077	if mmath == xr :
2078	utab = utab.assign_coords({ay:np.arange(ftab.shape[jy])+1.})
2079	if zdim and az and az != zdim :
2080	utab = utab.rename( {az:zdim})
2081	return utab
2082
2083	def f2v (ftab, nperio=None, psgn=1.0, zdim=None, action='ave') :
2084	'''Interpolates an array from F grid to V grid (i-mean)
2085	'''
2086	mmath = __mmath__ (ftab)
2087	ftab_0 = mmath.where ( np.isnan(ftab), 0., ftab)
2088	ftab_0 = lbc_add (ftab_0 , nperio=nperio, cd_type='F', psgn=psgn)
2089	ax, ix = __find_axis__ (ftab_0, 'x')
2090	az = __find_axis__ (ftab_0, 'z')[0]
2091	if ix :
2092	if action == 'ave' :
2093	vtab = 0.5 * (ftab_0 + np.roll (ftab_0, axis=ix, shift=-1))
2094	if action == 'min' :
2095	vtab = np.minimum (ftab_0 , np.roll (ftab_0, axis=ix, shift=-1))
2096	if action == 'max' :
2097	vtab = np.maximum (ftab_0 , np.roll (ftab_0, axis=ix, shift=-1))
2098	if action == 'mult' :
2099	vtab = ftab_0 * np.roll (ftab_0, axis=ix, shift=-1)
2100	vtab = lbc_del (vtab , nperio=nperio, cd_type='V', psgn=psgn)
2101	else :
2102	vtab = lbc_del (ftab_0, nperio=nperio, cd_type='V', psgn=psgn)
2103
2104	if mmath == xr :
2105	vtab = vtab.assign_coords({ax:np.arange(ftab.shape[ix])+1.})
2106	if zdim and az :
2107	if az != zdim :
2108	vtab = vtab.rename( {az:zdim})
2109	return vtab
2110
2111	def w2t (wtab, zcoord=None, zdim=None, sval=np.nan) :
2112	'''Interpolates an array on W grid to T grid (k-mean)
2113
2114	sval is the bottom value
2115	'''
2116	mmath = __mmath__ (wtab)
2117	wtab_0 = mmath.where ( np.isnan(wtab), 0., wtab)
2118
2119	az, kz = __find_axis__ (wtab_0, 'z')
2120
2121	if kz :
2122	ttab = 0.5 * ( wtab_0 + np.roll (wtab_0, axis=kz, shift=-1) )
2123	else :
2124	ttab = wtab_0
2125
2126	if mmath == xr :
2127	ttab[{az:kz}] = sval
2128	if zdim and az :
2129	if az != zdim :
2130	ttab = ttab.rename ( {az:zdim} )
2131	if zcoord is not None :
2132	ttab = ttab.assign_coords ( {zdim:zcoord} )
2133	else :
2134	ttab[..., -1, :, :] = sval
2135
2136	return ttab
2137
2138	def t2w (ttab, zcoord=None, zdim=None, sval=np.nan, extrap_surf=False) :
2139	'''Interpolates an array from T grid to W grid (k-mean)
2140
2141	sval is the surface value
2142	if extrap_surf==True, surface value is taken from 1st level value.
2143	'''
2144	mmath = __mmath__ (ttab)
2145	ttab_0 = mmath.where ( np.isnan(ttab), 0., ttab)
2146	az, kz = __find_axis__ (ttab_0, 'z')
2147	wtab = 0.5 * ( ttab_0 + np.roll (ttab_0, axis=kz, shift=1) )
2148
2149	if mmath == xr :
2150	if extrap_surf :
2151	wtab[{az:0}] = ttab[{az:0}]
2152	else :
2153	wtab[{az:0}] = sval
2154	else :
2155	if extrap_surf :
2156	wtab[..., 0, :, :] = ttab[..., 0, :, :]
2157	else :
2158	wtab[..., 0, :, :] = sval
2159
2160	if mmath == xr :
2161	if zdim and az and az != zdim :
2162	wtab = wtab.rename ( {az:zdim})
2163	if zcoord is not None :
2164	wtab = wtab.assign_coords ( {zdim:zcoord})
2165	else :
2166	wtab = wtab.assign_coords ( {zdim:np.arange(ttab.shape[kz])+1.} )
2167	return wtab
2168
2169	def fill (ptab, nperio, cd_type='T', npass=1, sval=np.nan) :
2170	'''Fills np.nan values with mean of neighbours
2171
2172	Inputs :
2173	ptab : input field to fill
2174	nperio, cd_type : periodicity characteristics
2175	'''
2176
2177	mmath = __mmath__ (ptab)
2178
2179	do_perio = False
2180	lperio = nperio
2181	if nperio == 4.2 :
2182	do_perio, lperio = True, 4
2183	if nperio == 6.2 :
2184	do_perio, lperio = True, 6
2185
2186	if do_perio :
2187	ztab = lbc_add (ptab, nperio=nperio)
2188	else :
2189	ztab = ptab
2190
2191	if np.isnan (sval) :
2192	ztab = mmath.where (np.isnan(ztab), np.nan, ztab)
2193	else :
2194	ztab = mmath.where (ztab==sval , np.nan, ztab)
2195
2196	for _ in np.arange (npass) :
2197	zmask = mmath.where ( np.isnan(ztab), 0., 1. )
2198	ztab0 = mmath.where ( np.isnan(ztab), 0., ztab )
2199	# Compte du nombre de voisins
2200	zcount = 1./6. * ( zmask \
2201	+ np.roll(zmask, shift=1, axis=-1) + np.roll(zmask, shift=-1, axis=-1) \
2202	+ np.roll(zmask, shift=1, axis=-2) + np.roll(zmask, shift=-1, axis=-2) \
2203	+ 0.5 * ( \
2204	+ np.roll(np.roll(zmask, shift= 1, axis=-2), shift= 1, axis=-1) \
2205	+ np.roll(np.roll(zmask, shift=-1, axis=-2), shift= 1, axis=-1) \
2206	+ np.roll(np.roll(zmask, shift= 1, axis=-2), shift=-1, axis=-1) \
2207	+ np.roll(np.roll(zmask, shift=-1, axis=-2), shift=-1, axis=-1) ) )
2208
2209	znew =1./6. * ( ztab0 \
2210	+ np.roll(ztab0, shift=1, axis=-1) + np.roll(ztab0, shift=-1, axis=-1) \
2211	+ np.roll(ztab0, shift=1, axis=-2) + np.roll(ztab0, shift=-1, axis=-2) \
2212	+ 0.5 * ( \
2213	+ np.roll(np.roll(ztab0 , shift= 1, axis=-2), shift= 1, axis=-1) \
2214	+ np.roll(np.roll(ztab0 , shift=-1, axis=-2), shift= 1, axis=-1) \
2215	+ np.roll(np.roll(ztab0 , shift= 1, axis=-2), shift=-1, axis=-1) \
2216	+ np.roll(np.roll(ztab0 , shift=-1, axis=-2), shift=-1, axis=-1) ) )
2217
2218	zcount = lbc (zcount, nperio=lperio, cd_type=cd_type)
2219	znew = lbc (znew , nperio=lperio, cd_type=cd_type)
2220
2221	ztab = mmath.where (np.logical_and (zmask==0., zcount>0), znew/zcount, ztab)
2222
2223	ztab = mmath.where (zcount==0, sval, ztab)
2224	if do_perio :
2225	ztab = lbc_del (ztab, nperio=lperio)
2226
2227	return ztab
2228
2229	def correct_uv (u, v, lat) :
2230	'''
2231	Corrects a Cartopy bug in orthographic projection
2232
2233	See https://github.com/SciTools/cartopy/issues/1179
2234
2235	The correction is needed with cartopy <= 0.20
2236	It seems that version 0.21 will correct the bug (https://github.com/SciTools/cartopy/pull/1926)
2237	Later note : the bug is still present in Cartopy 0.22
2238
2239	Inputs :
2240	u, v : eastward/northward components
2241	lat : latitude of the point (degrees north)
2242
2243	Outputs :
2244	modified eastward/nothward components to have correct polar projections in cartopy
2245	'''
2246	uv = np.sqrt (uu + vv) # Original modulus
2247	zu = u
2248	zv = v * np.cos (RAD*lat)
2249	zz = np.sqrt ( zuzu + zvzv ) # Corrected modulus
2250	uc = zu*uv/zz
2251	vc = zv*uv/zz # Final corrected values
2252	return uc, vc
2253
2254	def norm_uv (u, v) :
2255	'''Returns norm of a 2 components vector
2256	'''
2257	return np.sqrt (uu + vv)
2258
2259	def normalize_uv (u, v) :
2260	'''Normalizes 2 components vector
2261	'''
2262	uv = norm_uv (u, v)
2263	return u/uv, v/uv
2264
2265	def msf (vv, e1v_e3v, plat1d, depthw) :
2266	'''Computes the meridonal stream function
2267
2268	vv : meridional_velocity
2269	e1v_e3v : prodcut of scale factors e1v*e3v
2270	'''
2271
2272	v_e1v_e3v = vv * e1v_e3v
2273	v_e1v_e3v.attrs = vv.attrs
2274
2275	ax = __find_axis__ (v_e1v_e3v, 'x')[0]
2276	az = __find_axis__ (v_e1v_e3v, 'z')[0]
2277	if az == 'olevel' :
2278	new_az = 'olevel'
2279	else :
2280	new_az = 'depthw'
2281
2282	zomsf = -v_e1v_e3v.cumsum ( dim=az, keep_attrs=True).sum (dim=ax, keep_attrs=True)*1.E-6
2283	zomsf = zomsf - zomsf.isel ( { az:-1} )
2284
2285	ay = __find_axis__ (zomsf, 'y' )[0]
2286	zomsf = zomsf.assign_coords ( {az:depthw.values, ay:plat1d.values})
2287
2288	zomsf = zomsf.rename ( {ay:'lat'})
2289	if az != new_az :
2290	zomsf = zomsf.rename ( {az:new_az} )
2291	zomsf.attrs['standard_name'] = 'Meridional stream function'
2292	zomsf.attrs['long_name'] = 'Meridional stream function'
2293	zomsf.attrs['units'] = 'Sv'
2294	zomsf[new_az].attrs = depthw.attrs
2295	zomsf.lat.attrs=plat1d.attrs
2296
2297	return zomsf
2298
2299	def bsf (uu, e2u_e3u, mask, nperio=None, bsf0=None ) :
2300	'''Computes the barotropic stream function
2301
2302	uu : zonal_velocity
2303	e2u_e3u : product of scales factor e2u*e3u
2304	bsf0 : the point with bsf=0
2305	(ex: bsf0={'x':3, 'y':120} for orca2,
2306	bsf0={'x':5, 'y':300} for eORCA1
2307	'''
2308	u_e2u_e3u = uu * e2u_e3u
2309	u_e2u_e3u.attrs = uu.attrs
2310
2311	ay = __find_axis__ (u_e2u_e3u, 'y')[0]
2312	az = __find_axis__ (u_e2u_e3u, 'z')[0]
2313
2314	zbsf = -u_e2u_e3u.cumsum ( dim=ay, keep_attrs=True )
2315	zbsf = zbsf.sum (dim=az, keep_attrs=True)*1.E-6
2316
2317	if bsf0 :
2318	zbsf = zbsf - zbsf.isel (bsf0)
2319
2320	zbsf = zbsf.where (mask !=0, np.nan)
2321	zbsf.attrs.update (uu.attrs)
2322	zbsf.attrs['standard_name'] = 'barotropic_stream_function'
2323	zbsf.attrs['long_name'] = 'Barotropic stream function'
2324	zbsf.attrs['units'] = 'Sv'
2325	zbsf = lbc (zbsf, nperio=nperio, cd_type='F')
2326
2327	return zbsf
2328
2329	if f90nml :
2330	def namelist_read (ref=None, cfg=None, out='dict', flat=False, verbose=False) :
2331	'''Read NEMO namelist(s) and return either a dictionnary or an xarray dataset
2332
2333	ref : file with reference namelist, or a f90nml.namelist.Namelist object
2334	cfg : file with config namelist, or a f90nml.namelist.Namelist object
2335	At least one namelist neaded
2336
2337	out:
2338	'dict' to return a dictonnary
2339	'xr' to return an xarray dataset
2340	flat : only for dict output. Output a flat dictionary with all values.
2341
2342	'''
2343	if ref :
2344	if isinstance (ref, str) :
2345	nml_ref = f90nml.read (ref)
2346	if isinstance (ref, f90nml.namelist.Namelist) :
2347	nml_ref = ref
2348
2349	if cfg :
2350	if isinstance (cfg, str) :
2351	nml_cfg = f90nml.read (cfg)
2352	if isinstance (cfg, f90nml.namelist.Namelist) :
2353	nml_cfg = cfg
2354
2355	if out == 'dict' :
2356	dict_namelist = {}
2357	if out == 'xr' :
2358	xr_namelist = xr.Dataset ()
2359
2360	list_nml = []
2361	list_comment = []
2362
2363	if ref :
2364	list_nml.append (nml_ref)
2365	list_comment.append ('ref')
2366	if cfg :
2367	list_nml.append (nml_cfg)
2368	list_comment.append ('cfg')
2369
2370	for nml, comment in zip (list_nml, list_comment) :
2371	if verbose :
2372	print (comment)
2373	if flat and out =='dict' :
2374	for nam in nml.keys () :
2375	if verbose :
2376	print (nam)
2377	for value in nml[nam] :
2378	if out == 'dict' :
2379	dict_namelist[value] = nml[nam][value]
2380	if verbose :
2381	print (nam, ':', value, ':', nml[nam][value])
2382	else :
2383	for nam in nml.keys () :
2384	if verbose :
2385	print (nam)
2386	if out == 'dict' :
2387	if nam not in dict_namelist.keys () :
2388	dict_namelist[nam] = {}
2389	for value in nml[nam] :
2390	if out == 'dict' :
2391	dict_namelist[nam][value] = nml[nam][value]
2392	if out == 'xr' :
2393	xr_namelist[value] = nml[nam][value]
2394	if verbose :
2395	print (nam, ':', value, ':', nml[nam][value])
2396
2397	if out == 'dict' :
2398	return dict_namelist
2399	if out == 'xr' :
2400	return xr_namelist
2401	else :
2402	def namelist_read (ref=None, cfg=None, out='dict', flat=False, verbose=False) :
2403	'''Shadow version of namelist read, when f90nm module was not found
2404
2405	namelist_read :
2406	Read NEMO namelist(s) and return either a dictionnary or an xarray dataset
2407	'''
2408	print ( 'Error : module f90nml not found' )
2409	print ( 'Cannot call namelist_read' )
2410	print ( 'Call parameters where : ')
2411	print ( f'{err=} {ref=} {cfg=} {out=} {flat=} {verbose=}' )
2412
2413	def fill_closed_seas (imask, nperio=None, cd_type='T') :
2414	'''Fill closed seas with image processing library
2415
2416	imask : mask, 1 on ocean, 0 on land
2417	'''
2418	from scipy import ndimage
2419
2420	imask_filled = ndimage.binary_fill_holes ( lbc (imask, nperio=nperio, cd_type=cd_type))
2421	imask_filled = lbc ( imask_filled, nperio=nperio, cd_type=cd_type)
2422
2423	return imask_filled
2424
2425	# ======================================================
2426	# Sea water state function parameters from NEMO code
2427
2428	RDELTAS = 32.
2429	R1_S0 = 0.875/35.16504
2430	R1_T0 = 1./40.
2431	R1_Z0 = 1.e-4
2432
2433	EOS000 = 8.0189615746e+02
2434	EOS100 = 8.6672408165e+02
2435	EOS200 = -1.7864682637e+03
2436	EOS300 = 2.0375295546e+03
2437	EOS400 = -1.2849161071e+03
2438	EOS500 = 4.3227585684e+02
2439	EOS600 = -6.0579916612e+01
2440	EOS010 = 2.6010145068e+01
2441	EOS110 = -6.5281885265e+01
2442	EOS210 = 8.1770425108e+01
2443	EOS310 = -5.6888046321e+01
2444	EOS410 = 1.7681814114e+01
2445	EOS510 = -1.9193502195
2446	EOS020 = -3.7074170417e+01
2447	EOS120 = 6.1548258127e+01
2448	EOS220 = -6.0362551501e+01
2449	EOS320 = 2.9130021253e+01
2450	EOS420 = -5.4723692739
2451	EOS030 = 2.1661789529e+01
2452	EOS130 = -3.3449108469e+01
2453	EOS230 = 1.9717078466e+01
2454	EOS330 = -3.1742946532
2455	EOS040 = -8.3627885467
2456	EOS140 = 1.1311538584e+01
2457	EOS240 = -5.3563304045
2458	EOS050 = 5.4048723791e-01
2459	EOS150 = 4.8169980163e-01
2460	EOS060 = -1.9083568888e-01
2461	EOS001 = 1.9681925209e+01
2462	EOS101 = -4.2549998214e+01
2463	EOS201 = 5.0774768218e+01
2464	EOS301 = -3.0938076334e+01
2465	EOS401 = 6.6051753097
2466	EOS011 = -1.3336301113e+01
2467	EOS111 = -4.4870114575
2468	EOS211 = 5.0042598061
2469	EOS311 = -6.5399043664e-01
2470	EOS021 = 6.7080479603
2471	EOS121 = 3.5063081279
2472	EOS221 = -1.8795372996
2473	EOS031 = -2.4649669534
2474	EOS131 = -5.5077101279e-01
2475	EOS041 = 5.5927935970e-01
2476	EOS002 = 2.0660924175
2477	EOS102 = -4.9527603989
2478	EOS202 = 2.5019633244
2479	EOS012 = 2.0564311499
2480	EOS112 = -2.1311365518e-01
2481	EOS022 = -1.2419983026
2482	EOS003 = -2.3342758797e-02
2483	EOS103 = -1.8507636718e-02
2484	EOS013 = 3.7969820455e-01
2485
2486	def rhop ( ptemp, psal ) :
2487	'''Returns potential density referenced to surface
2488
2489	Computation from NEMO code
2490	'''
2491	zt = ptemp * R1_T0 # Temperature (Â°C)
2492	zs = np.sqrt ( np.abs( psal + RDELTAS ) * R1_S0 ) # Square root of salinity (PSS)
2493	#
2494	prhop = (
2495	(((((EOS060*zt
2496	+ EOS150zs + EOS050)zt
2497	+ (EOS240zs + EOS140)zs + EOS040)*zt
2498	+ ((EOS330zs + EOS230)zs + EOS130)zs + EOS030)zt
2499	+ (((EOS420zs + EOS320)zs + EOS220)zs + EOS120)zs + EOS020)*zt
2500	+ ((((EOS510zs + EOS410)zs + EOS310)zs + EOS210)zs + EOS110)zs + EOS010)zt
2501	+ (((((EOS600zs+ EOS500)zs + EOS400)zs + EOS300)zs + EOS200)zs + EOS100)zs + EOS000 )
2502	#
2503	return prhop
2504
2505	def rho ( pdep, ptemp, psal ) :
2506	'''Returns in situ density
2507
2508	Computation from NEMO code
2509	'''
2510	zh = pdep * R1_Z0 # Depth (m)
2511	zt = ptemp * R1_T0 # Temperature (Â°C)
2512	zs = np.sqrt ( np.abs( psal + RDELTAS ) * R1_S0 ) # Square root salinity (PSS)
2513	#
2514	zn3 = EOS013zt + EOS103zs+EOS003
2515	#
2516	zn2 = (EOS022zt + EOS112zs+EOS012)zt + (EOS202zs+EOS102)*zs+EOS002
2517	#
2518	zn1 = (
2519	(((EOS041*zt
2520	+ EOS131zs + EOS031)zt
2521	+ (EOS221zs + EOS121)zs + EOS021)*zt
2522	+ ((EOS311zs + EOS211)zs + EOS111)zs + EOS011)zt
2523	+ (((EOS401zs + EOS301)zs + EOS201)zs + EOS101)zs + EOS001 )
2524	#
2525	zn0 = (
2526	(((((EOS060*zt
2527	+ EOS150zs + EOS050)zt
2528	+ (EOS240zs + EOS140)zs + EOS040)*zt
2529	+ ((EOS330zs + EOS230)zs + EOS130)zs + EOS030)zt
2530	+ (((EOS420zs + EOS320)zs + EOS220)zs + EOS120)zs + EOS020)*zt
2531	+ ((((EOS510zs + EOS410)zs + EOS310)zs + EOS210)zs + EOS110)zs + EOS010)zt
2532	+ (((((EOS600zs + EOS500)zs + EOS400)zs + EOS300)zs +
2533	EOS200)zs + EOS100)zs + EOS000 )
2534	#
2535	prho = ( ( zn3 * zh + zn2 ) * zh + zn1 ) * zh + zn0
2536	#
2537	return prho
2538
2539	## ===========================================================================
2540	##
2541	## That's all folk's !!!
2542	##
2543	## ===========================================================================
2544
2545	# def __is_orca_north_fold__ ( Xtest, cname_long='T' ) :
2546	# '''
2547	# Ported (pirated !!?) from Sosie
2548
2549	# Tell if there is a 2/point band overlaping folding at the north pole typical of the ORCA grid
2550
2551	# 0 => not an orca grid (or unknown one)
2552	# 4 => North fold T-point pivot (ex: ORCA2)
2553	# 6 => North fold F-point pivot (ex: ORCA1)
2554
2555	# We need all this 'cname_long' stuff because with our method, there is a
2556	# confusion between "Grid_U with T-fold" and "Grid_V with F-fold"
2557	# => so knowing the name of the longitude array (as in namelist, and hence as
2558	# in netcdf file) might help taking the righ decision !!! UGLY!!!
2559	# => not implemented yet
2560	# '''
2561
2562	# ifld_nord = 0 ; cgrd_type = 'X'
2563	# ny, nx = Xtest.shape[-2:]
2564
2565	# if ny > 3 : # (case if called with a 1D array, ignoring...)
2566	# if ( Xtest [ny-1, 1:nx//2-1] - Xtest [ny-3, nx-1:nx-nx//2+1:-1] ).sum() == 0. :
2567	# ifld_nord = 4 ; cgrd_type = 'T' # T-pivot, grid_T
2568
2569	# if ( Xtest [ny-1, 1:nx//2-1] - Xtest [ny-3, nx-2:nx-nx//2 :-1] ).sum() == 0. :
2570	# if cnlon == 'U' : ifld_nord = 4 ; cgrd_type = 'U' # T-pivot, grid_T
2571	# ## LOLO: PROBLEM == 6, V !!!
2572
2573	# if ( Xtest [ny-1, 1:nx//2-1] - Xtest [ny-3, nx-1:nx-nx//2+1:-1] ).sum() == 0. :
2574	# ifld_nord = 4 ; cgrd_type = 'V' # T-pivot, grid_V
2575
2576	# if ( Xtest [ny-1, 1:nx//2-1] - Xtest [ny-2, nx-1-1:nx-nx//2:-1] ).sum() == 0. :
2577	# ifld_nord = 6 ; cgrd_type = 'T'# F-pivot, grid_T
2578
2579	# if ( Xtest [ny-1, 1:nx//2-1] - Xtest [ny-1, nx-1:nx-nx//2-1:-1] ).sum() == 0. :
2580	# ifld_nord = 6 ; cgrd_type = 'U' # F-pivot, grid_U
2581
2582	# if ( Xtest [ny-1, 1:nx//2-1] - Xtest [ny-3, nx-2:nx-nx//2 :-1] ).sum() == 0. :
2583	# if cnlon == 'V' : ifld_nord = 6 ; cgrd_type = 'V' # F-pivot, grid_V
2584	# ## LOLO: PROBLEM == 4, U !!!
2585
2586	# return ifld_nord, cgrd_type

Note: See TracBrowser for help on using the repository browser.

Context Navigation

source: TOOLS/WATER_BUDGET/nemo.py @ 6709

Download in other formats: