[5948] | 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 | ## =========================================================================== |
---|