1 | # -*- coding: utf-8 -*- |
---|
2 | ## =========================================================================== |
---|
3 | ## |
---|
4 | ## Warning, to install, configure, run, use any of Olivier Marti's |
---|
5 | ## software or to read the associated documentation you'll need at least |
---|
6 | ## one (1) brain in a reasonably working order. Lack of this implement |
---|
7 | ## will void any warranties (either express or implied). |
---|
8 | ## O. Marti assumes no responsability for errors, omissions, |
---|
9 | ## data loss, or any other consequences caused directly or indirectly by |
---|
10 | ## the usage of his software by incorrectly or partially configured |
---|
11 | ## personal. |
---|
12 | ## |
---|
13 | ''' |
---|
14 | Utilities to plot NEMO ORCA fields |
---|
15 | Periodicity and other stuff |
---|
16 | |
---|
17 | olivier.marti@lsce.ipsl.fr |
---|
18 | ''' |
---|
19 | |
---|
20 | import sys |
---|
21 | |
---|
22 | try : import numpy as np |
---|
23 | except : pass |
---|
24 | |
---|
25 | try : import xarray as xr |
---|
26 | except : pass |
---|
27 | |
---|
28 | ## SVN information |
---|
29 | __Author__ = "$Author$" |
---|
30 | __Date__ = "$Date$" |
---|
31 | __Revision__ = "$Revision$" |
---|
32 | __Id__ = "$Id$" |
---|
33 | __HeadURL = "$HeadURL$" |
---|
34 | |
---|
35 | def __guessNperio__ (jpi, nperio=None) : |
---|
36 | ''' |
---|
37 | Tries to guess the value of nperio (periodicity parameter. See NEMO documentation for details) |
---|
38 | |
---|
39 | Inputs |
---|
40 | jpi : number of longitudes |
---|
41 | nperio : periodicity parameter |
---|
42 | ''' |
---|
43 | |
---|
44 | if nperio == None : |
---|
45 | if jpi == 182 : nperio = 4 # ORCA2. We choose legacy orca2. |
---|
46 | if jpi == 362 : nperio = 6 # ORCA1. |
---|
47 | if jpi == 1442 : nperio = 6 # ORCA025. |
---|
48 | # |
---|
49 | if nperio == None : |
---|
50 | sys.exit ('in nemo.lbc : nperio not found, and cannot by guessed' ) |
---|
51 | else : |
---|
52 | print ('nperio set as {:d} (deduced from jpi={:d})'.format(nperio, jpi)) |
---|
53 | |
---|
54 | return nperio |
---|
55 | |
---|
56 | def fixed_lon (lon) : |
---|
57 | ''' |
---|
58 | Returns corrected longitudes for nicer plots |
---|
59 | |
---|
60 | fixed_lon (lon) |
---|
61 | lon : longitudes of the grid. At least 2D. |
---|
62 | ''' |
---|
63 | fixed_lon = lon.copy () |
---|
64 | for i, start in enumerate (np.argmax (np.abs (np.diff (lon, axis=-1)) > 180, axis=-1)) : |
---|
65 | fixed_lon[...,i, start+1:] += 360 |
---|
66 | return fixed_lon |
---|
67 | |
---|
68 | def jeq (lat) : |
---|
69 | ''' |
---|
70 | Returns j index of equator in the grid |
---|
71 | |
---|
72 | lat : latitudes of the grid. At least 2D. |
---|
73 | ''' |
---|
74 | jeq = np.argmin (np.abs (np.float64(lat[:,0]))) |
---|
75 | return jeq |
---|
76 | |
---|
77 | def lon1D (lon, lat=None) : |
---|
78 | ''' |
---|
79 | Returns 1D longitude for simple plots. |
---|
80 | |
---|
81 | lon : longitudes of the grid |
---|
82 | lat (optionnal) : latitudes of the grid |
---|
83 | ''' |
---|
84 | |
---|
85 | if np.max (lat) != None : |
---|
86 | je = jeq (lat) |
---|
87 | lon1D = lon[..., je, :] |
---|
88 | else : |
---|
89 | jpj, jpi = lon.shape[-2:] |
---|
90 | lon1D = lon[..., jpj//3, :] |
---|
91 | |
---|
92 | start = np.argmax (np.abs (np.diff (lon1D, axis=-1)) > 180.0, axis=-1) |
---|
93 | lon1D[..., start+1:] += 360 |
---|
94 | |
---|
95 | return lon1D |
---|
96 | |
---|
97 | def latreg (lat) : |
---|
98 | ''' |
---|
99 | Returns maximum j index wehre gridlines are along latitude in the norethern hemisphere |
---|
100 | |
---|
101 | lon : longitudes of the grid |
---|
102 | lat (optionnal) : latitudes of the grid |
---|
103 | ''' |
---|
104 | dy = np.float64 (np.mean (np.abs (lat - np.roll(lat,shift=1,axis=-2)))) |
---|
105 | je = jeq (lat) |
---|
106 | jreg = np.where (lat[...,je:,:].max(axis=-1) - lat[...,je:,:].min(axis=-1)< dy/100.)[-1][-1] + je |
---|
107 | latreg = np.float64 (lat[...,jreg,:].mean(axis=-1)) |
---|
108 | JREG = jreg |
---|
109 | |
---|
110 | return jreg, latreg |
---|
111 | |
---|
112 | def lat1D (lat) : |
---|
113 | ''' |
---|
114 | Returns 1D latitudes for zonal means and simple plots. |
---|
115 | |
---|
116 | lat : latitudes of the grid |
---|
117 | ''' |
---|
118 | jpj, jpi = lat.shape[-2:] |
---|
119 | try : |
---|
120 | if type (lat) == xr.core.dataarray.DataArray : math = xr |
---|
121 | else : math = np |
---|
122 | except : math = np |
---|
123 | |
---|
124 | # jreg : index of last uniform latitude, north of equator |
---|
125 | |
---|
126 | dy = np.float64 (np.mean (np.abs (lat - np.roll (lat, shift=1,axis=-2)))) |
---|
127 | je = jeq (lat) |
---|
128 | lateq = np.float64 (lat[...,je,:].mean(axis=-1)) |
---|
129 | |
---|
130 | jreg = np.where (lat[...,je:,:].max(axis=-1) - lat[...,je:,:].min(axis=-1)< dy/100.)[-1][-1] + je |
---|
131 | latreg = np.float64 (lat[...,jreg,:].mean(axis=-1)) |
---|
132 | latave = np.mean (lat, axis=-1) |
---|
133 | |
---|
134 | if ( np.abs (lateq) < dy/100. ) : # T, U or W grid |
---|
135 | dys = (90.-latreg)//(jpj-jreg-1)*0.5 |
---|
136 | yrange = 90.-dys-latreg |
---|
137 | else : # V or F grid |
---|
138 | yrange = 90. -latreg |
---|
139 | |
---|
140 | lat1D = math.where (latave<latreg, latave, latreg + yrange * (np.arange(jpj)-jreg)/(jpj-jreg-1) ) |
---|
141 | |
---|
142 | if math == xr : |
---|
143 | lat1D.attrs = lat.attrs |
---|
144 | |
---|
145 | return lat1D |
---|
146 | |
---|
147 | def latlon1D (lat, lon) : |
---|
148 | ''' |
---|
149 | Returns simple latitude and longitude (1D) for simple plots. |
---|
150 | ''' |
---|
151 | laty = lat1D (lat) |
---|
152 | lonx = lon1D (lon, lat) |
---|
153 | |
---|
154 | return laty, lonx |
---|
155 | |
---|
156 | def extend (tab, Lon=False, jplus=25) : |
---|
157 | ''' |
---|
158 | Returns extended field eastward to have better plots |
---|
159 | Works only for xarray and numpy data (?) |
---|
160 | |
---|
161 | tab : field to extend. |
---|
162 | Lon : (optional, default=False) : if True, add 360 in the extended parts of the field |
---|
163 | jplus (optional, default=25) : number of points added on the east side of the field |
---|
164 | |
---|
165 | ''' |
---|
166 | jpi = tab.shape[-1] |
---|
167 | JPLUS = jplus |
---|
168 | |
---|
169 | try : |
---|
170 | if type (tab) == xr.core.dataarray.DataArray : math = xr |
---|
171 | else : math = np |
---|
172 | except : math = np |
---|
173 | |
---|
174 | if Lon : xplus = -360.0 |
---|
175 | else : xplus = 0.0 |
---|
176 | |
---|
177 | if math == xr : |
---|
178 | extend = xr.concat ( (tab[..., 2:jpi] + xplus, tab[..., 2:jplus]), dim=tab.dims[-1] ) |
---|
179 | if math == np : |
---|
180 | extend = np.concatenate ( (tab[..., 2:jpi] + xplus, tab[..., 2:jplus]), axis=-1 ) |
---|
181 | |
---|
182 | return extend |
---|
183 | |
---|
184 | |
---|
185 | def lbc (ptab, nperio=None, cd_type='T', psgn=1.0, nemo_4U_bug=False) : |
---|
186 | """ |
---|
187 | Set periodicity on input field |
---|
188 | ptab : Input array (works |
---|
189 | rank 2 at least : ptab[...., lat, lon] |
---|
190 | nperio : Type of periodicity |
---|
191 | 1, 4, 6 : Cyclic on i dimension (generaly longitudes) |
---|
192 | 2 : Obsolete (was symmetric condition at southern boundary ?) |
---|
193 | 3, 4 : North fold T-point pivot (legacy ORCA2) |
---|
194 | 5, 6 : North fold F-point pivot (ORCA1, ORCA025, ORCA2 with new grid for paleo) |
---|
195 | cd_type : Grid specification : T, U, V or F |
---|
196 | psgn : For change of sign for vector components (1 for scalars, -1 for vector components) |
---|
197 | |
---|
198 | See NEMO documentation for further details |
---|
199 | """ |
---|
200 | |
---|
201 | jpi = ptab.shape[-1] |
---|
202 | nperio = __guessNperio__ ( jpi, nperio) |
---|
203 | psgn = ptab.dtype.type(psgn) |
---|
204 | |
---|
205 | # |
---|
206 | #> East-West boundary conditions |
---|
207 | # ------------------------------ |
---|
208 | if nperio in [1, 4, 6] : |
---|
209 | # ... cyclic |
---|
210 | ptab [..., :, 0] = ptab [..., :, -2] |
---|
211 | ptab [..., :, -1] = ptab [..., :, 1] |
---|
212 | |
---|
213 | # |
---|
214 | #> North-South boundary conditions |
---|
215 | # -------------------------------- |
---|
216 | if nperio in [3, 4] : # North fold T-point pivot |
---|
217 | if cd_type in [ 'T', 'W' ] : # T-, W-point |
---|
218 | ptab [..., -1, 1: ] = psgn * ptab [..., -3, -1:0:-1 ] |
---|
219 | ptab [..., -1, 0 ] = psgn * ptab [..., -3, 2 ] |
---|
220 | ptab [..., -2, jpi//2: ] = psgn * ptab [..., -2, jpi//2:0:-1 ] |
---|
221 | |
---|
222 | if cd_type == 'U' : |
---|
223 | ptab [..., -1, 0:-1 ] = psgn * ptab [..., -3, -1:0:-1 ] |
---|
224 | ptab [..., -1, 0 ] = psgn * ptab [..., -3, 1 ] |
---|
225 | ptab [..., -1, -1 ] = psgn * ptab [..., -3, -2 ] |
---|
226 | |
---|
227 | if nemo_4U_bug : |
---|
228 | ptab [..., -2, jpi//2+1:-1] = psgn * ptab [..., -2, jpi//2-2:0:-1] |
---|
229 | ptab [..., -2, jpi//2-1 ] = psgn * ptab [..., -2, jpi//2 ] |
---|
230 | else : |
---|
231 | ptab [..., -2, jpi//2-1:-1] = psgn * ptab [..., -2, jpi//2:0:-1] |
---|
232 | |
---|
233 | if cd_type == 'V' : |
---|
234 | ptab [..., -2, 1: ] = psgn * ptab [..., -3, jpi-1:0:-1 ] |
---|
235 | ptab [..., -1, 1: ] = psgn * ptab [..., -4, -1:0:-1 ] |
---|
236 | ptab [..., -1, 0 ] = psgn * ptab [..., -4, 2 ] |
---|
237 | |
---|
238 | if cd_type == 'F' : |
---|
239 | ptab [..., -2, 0:-1 ] = psgn * ptab [..., -3, -1:0:-1 ] |
---|
240 | ptab [..., -1, 0:-1 ] = psgn * ptab [..., -4, -1:0:-1 ] |
---|
241 | ptab [..., -1, 0 ] = psgn * ptab [..., -4, 1 ] |
---|
242 | ptab [..., -1, -1 ] = psgn * ptab [..., -4, -2 ] |
---|
243 | |
---|
244 | if nperio in [5, 6] : # North fold F-point pivot |
---|
245 | if cd_type in ['T', 'W'] : |
---|
246 | ptab [..., -1, 0: ] = psgn * ptab [..., -2, -1::-1 ] |
---|
247 | |
---|
248 | if cd_type == 'U' : |
---|
249 | ptab [..., -1, 0:-1 ] = psgn * ptab [..., -2, -2::-1 ] |
---|
250 | ptab [..., -1, -1 ] = psgn * ptab [..., -2, 0 ] |
---|
251 | |
---|
252 | if cd_type == 'V' : |
---|
253 | ptab [..., -1, 0: ] = psgn * ptab [..., -3, -1::-1 ] |
---|
254 | ptab [..., -2, jpi//2: ] = psgn * ptab [..., -2, jpi//2-1::-1 ] |
---|
255 | |
---|
256 | if cd_type == 'F' : |
---|
257 | ptab [..., -1, 0:-1 ] = psgn * ptab [..., -3, -2::-1 ] |
---|
258 | ptab [..., -1, -1 ] = psgn * ptab [..., -3, 0 ] |
---|
259 | ptab [..., -2, jpi//2:-1] = psgn * ptab [..., -2, jpi//2-2::-1 ] |
---|
260 | |
---|
261 | return ptab |
---|
262 | |
---|
263 | def lbc_mask (ptab, nperio=None, cd_type='T', sval=None) : |
---|
264 | # |
---|
265 | """ |
---|
266 | Mask fields on duplicated points |
---|
267 | ptab : Input array |
---|
268 | rank 2 at least : ptab[...., lat, lon] |
---|
269 | nperio : Type of periodicity |
---|
270 | 1, 4, 6 : Cyclic on i dimension (generaly longitudes) |
---|
271 | 2 : Obsolete (was symmetric condition at southern boundary ?) |
---|
272 | 3, 4 : North fold T-point pivot (legacy ORCA2) |
---|
273 | 5, 6 : North fold F-point pivot (ORCA1, ORCA025, ORCA2 with new grid for paleo) |
---|
274 | cd_type : Grid specification : T, U, V or F |
---|
275 | |
---|
276 | See NEMO documentation for further details |
---|
277 | """ |
---|
278 | |
---|
279 | jpi = ptab.shape[-1] |
---|
280 | nperio = __guessNperio__ (jpi, nperio) |
---|
281 | if sval == None : sval = ptab.dtype.type (0) |
---|
282 | |
---|
283 | # |
---|
284 | #> East-West boundary conditions |
---|
285 | # ------------------------------ |
---|
286 | if nperio in [1, 4, 6] : |
---|
287 | # ... cyclic |
---|
288 | ptab [...,:, 0] = sval |
---|
289 | ptab [...,:, -1] = sval |
---|
290 | |
---|
291 | # |
---|
292 | #> South (in which nperio cases ?) |
---|
293 | # -------------------------------- |
---|
294 | if nperio in [1, 3, 4, 5, 6] : |
---|
295 | ptab [...,0,:] = sval |
---|
296 | |
---|
297 | # |
---|
298 | #> North-South boundary conditions |
---|
299 | # -------------------------------- |
---|
300 | if nperio in [3, 4] : # North fold T-point pivot |
---|
301 | if cd_type in [ 'T', 'W' ] : # T-, W-point |
---|
302 | ptab [..., -1, 1: ] = sval |
---|
303 | ptab [..., -1, 0 ] = sval |
---|
304 | ptab [..., -2, jpi//2: ] = sval |
---|
305 | |
---|
306 | if cd_type == 'U' : |
---|
307 | ptab [..., -1, 0:-1 ] = sval |
---|
308 | ptab [..., -1, 0 ] = sval |
---|
309 | ptab [..., -1, -1 ] = sval |
---|
310 | ptab [..., -2, jpi//2 :] = sval |
---|
311 | |
---|
312 | if cd_type == 'V' : |
---|
313 | ptab [..., -2, 1: ] = sval |
---|
314 | ptab [..., -1, 1: ] = sval |
---|
315 | ptab [..., -1, 0 ] = sval |
---|
316 | |
---|
317 | if cd_type == 'F' : |
---|
318 | ptab [..., -2, 0:-1 ] = sval |
---|
319 | ptab [..., -1, 0:-1 ] = sval |
---|
320 | ptab [..., -1, 0 ] = sval |
---|
321 | ptab [..., -1, -1 ] = sval |
---|
322 | |
---|
323 | if nperio in [5, 6] : # North fold F-point pivot |
---|
324 | if cd_type in ['T', 'W'] : |
---|
325 | ptab [..., -1, 0: ] = sval |
---|
326 | |
---|
327 | if cd_type == 'U' : |
---|
328 | ptab [..., -1, 0:-1 ] = sval |
---|
329 | ptab [..., -1, -1 ] = sval |
---|
330 | |
---|
331 | if cd_type == 'V' : |
---|
332 | ptab [..., -1, 0: ] = sval |
---|
333 | ptab [..., -2, jpi//2: ] = sval |
---|
334 | |
---|
335 | if cd_type == 'F' : |
---|
336 | ptab [..., -1, 0:-1 ] = sval |
---|
337 | ptab [..., -1, -1 ] = sval |
---|
338 | ptab [..., -2, jpi//2:-1] = sval |
---|
339 | |
---|
340 | return ptab |
---|
341 | |
---|
342 | ## =========================================================================== |
---|
343 | ## |
---|
344 | ## That's all folk's !!! |
---|
345 | ## |
---|
346 | ## =========================================================================== |
---|