source: TOOLS/WATER_BUDGET/lmdz.py

Last change on this file was 6764, checked in by omamce, 4 months ago

O.M. : water budget - version mars 2024

  • Property svn:keywords set to Date Revision HeadURL Author
File size: 14.2 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'''
21Utilities for LMDZ grid
22
23- Lots of tests for xarray object
24- Not much tested for numpy objects
25
26Author: olivier.marti@lsce.ipsl.fr
27
28## SVN information
29__Author__   = "$Author$"
30__Date__     = "$Date$"
31__Revision__ = "$Revision$"
32__Id__       = "$Id: $"
33__HeadURL    = "$HeadURL$"
34'''
35
36import numpy as np
37import xarray as xr
38
39RPI   = np.pi
40RAD   = np.deg2rad (1.0)
41DAR   = np.rad2deg (1.0)
42REPSI = np.finfo (1.0).eps
43
44RAAMO    =      12          # Number of months in one year
45RJJHH    =      24          # Number of hours in one day
46RHHMM    =      60          # Number of minutes in one hour
47RMMSS    =      60          # Number of seconds in one minute
48RA       = 6371229.0        # Earth radius                                  [m]
49GRAV     =       9.80665    # Gravity                                       [m/s2]
50RT0      =     273.15       # Freezing point of fresh water                 [Kelvin]
51RAU0     =    1026.0        # Volumic mass of sea water                     [kg/m3]
52RLEVAP   =       2.5e+6     # Latent heat of evaporation (water)            [J/K]
53VKARMN   =       0.4        # Von Karman constant
54STEFAN   =       5.67e-8    # Stefan-Boltzmann constant                     [W/m2/K4]
55#RHOS     =     330.         # Volumic mass of snow                          [kg/m3]
56
57RDAY     = RJJHH * RHHMM * RMMSS                # Day length               [s]
58RSIYEA   = 365.25 * RDAY * 2. * RPI / 6.283076  # Sideral year length      [s]
59RSIDAY   = RDAY / (1. + RDAY / RSIYEA)          # Sideral day length       [s]
60ROMEGA   = 2. * RPI / RSIDAY                    # Earth rotation parameter [s-1]
61
62def __mmath__ (ptab, default=None) :
63    '''Determines the type of tab : xarray, numpy or numpy.ma object ?
64
65    Returns type
66    '''
67    mmath = default
68    if isinstance (ptab, xr.core.dataarray.DataArray) :
69        mmath = xr
70    if isinstance (ptab, np.ndarray) :
71        mmath = np
72    if isinstance (ptab, np.ma.MaskType) :
73        mmath = np.ma
74
75    return mmath
76
77#
78def extend (tab, Lon=False, jplus=25, jpi=None, lonplus=360.0) :
79    '''Returns extended field eastward to have better plots, and box average crossing the boundary
80
81    Works only for xarray and numpy data (?)
82
83    tab : field to extend.
84    Lon : (optional, default=False) : if True, add 360 in the extended parts of the field
85    jpi : normal longitude dimension of the field. exrtend does nothing it the actual
86        size of the field != jpi (avoid to extend several times)
87    jplus (optional, default=25) : number of points added on the east side of the field
88    '''
89   
90    math = __mmath__ (tab)
91    if tab.shape[-1] == 1 :
92        ztab = tab
93
94    else :
95        if jpi is None :
96            jpi = tab.shape[-1]
97
98        if Lon :
99            xplus =  lonplus
100        else   :
101            xplus =    0.0
102
103        if tab.shape[-1] > jpi :
104            ztab = tab
105        else :
106            istart = 0
107            le     = jpi+1
108            la     = 0
109            if math == xr :
110                lon = tab.dims[-1]
111                ztab         = xr.concat      ((
112                    tab.isel (lon=slice(istart   ,istart+le      )),
113                    tab.isel (lon=slice(istart+la,istart+la+jplus))+xplus  ), dim=lon)
114                if 'lon' in tab.dims :
115                    extend_lon = xr.concat      ((
116                        tab.coords[lon].isel(lon=slice(istart   ,istart+le      )),
117                        tab.coords[lon].isel(lon=slice(istart+la,istart+la+jplus))+lonplus), dim=lon)
118                    ztab = ztab.assign_coords ( {tab.dims[-1]:extend_lon} )
119                #except :
120                #    pass
121            if math == np :
122                ztab = np.concatenate ((tab    [..., istart:istart+le],
123                                        tab    [..., istart+la:jplus]+xplus  ), axis=-1)
124
125    return ztab
126
127def interp1d (x, xp, yp, zdim='presnivs', units=None, verbose=False, method='linear') :
128    '''One-dimensionnal interpolation of a multi-dimensionnal field
129
130    Intended to interpolate on standard pressure level
131
132    All inputs shoud be xarray data arrays
133
134    Input :
135       x      : levels at wich we want to interpolate (i.e. standard pressure levels)
136       xp     : position of the input points (i.e. pressure)
137       yp     : fields values at these points (temperature, humidity, etc ..)
138       zdim   : name of the dimension that we want to interpolate
139       method : available methods are
140           linear
141           log[arithmic]. Available for positive values only
142           nearest : nearest neighbor
143
144           Warning : xp should be decreasing values along zdim axis
145    '''
146    # Get the number of dimension with dim==zdim
147    axis = list(xp.dims).index(zdim)
148
149    # Get the number of levels in each arrays
150    nk_ou = len (x)
151
152    # Define the result array
153    in_shape       = np.array (xp.shape)
154    if verbose :
155        print ( f'{in_shape=}' )
156    ou_shape       = np.array (in_shape)
157    if verbose :
158        print ( f'{ou_shape=}' )
159    ou_shape[axis] = nk_ou
160
161    in_dims        = list (yp.dims)
162    if verbose :
163        print ( f'{in_dims=}' )
164    ou_dims        = in_dims
165
166    pdim           = x.dims[0]
167    ou_dims[axis]  = pdim
168
169    new_coords = []
170    for i, dim in enumerate (yp.dims) :
171        if dim == zdim :
172            ou_dims[i] = x.dims[0]
173            if units is not None :
174                yp[dim].attrs['units'] = units
175            new_coords.append (x             .values)
176        else :
177            new_coords.append (yp.coords[dim].values)
178
179    if verbose :
180        print ( f'{ou_dims   =}' )
181        print ( f'{new_coords=}' )
182
183    ou_tab = xr.DataArray (np.empty (ou_shape), dims=ou_dims, coords=new_coords)
184
185    if 'log' in method :
186        yp_min = yp.min()
187        yp_max = yp.max()
188        indic  = yp_min * yp_max
189        if indic <= 0. :
190            print ( 'Input data have a change of sign')
191            print ( 'Error: logarithmic method is available only for')
192            print ( 'positive or negative input values. ')
193            raise ValueError
194
195    # Optimized (pre-compiled) interpolation loop
196    #@numba.jit (nopython=True)
197    def __interp (x, xp, yp) :
198        # Interpolate
199        # Find index of the just above level
200        idk1 = np.minimum ( (x-xp), 0.).argmax (dim=zdim)
201        idk2 = idk1 - 1
202        idk2 = np.maximum (idk2, 0)
203
204        x1   = xp[{zdim:idk1}]
205        x2   = xp[{zdim:idk2}]
206
207        dx1  = x  - x1
208        dx2  = x2 - x
209        dx   = x2 - x1
210        dx1  = dx1/dx
211        dx2  = dx2/dx
212
213        y1   = yp[{zdim:idk1}]
214        y2   = yp[{zdim:idk2}]
215
216        if 'linear'  in method :
217            result = dx1*y2 + dx2*y1
218        if 'log'     in method :
219            if yp_min > 0. :
220                result = np.power(y2, dx1) * np.power(y1, dx2)
221            if yp_max < 0. :
222                result = -np.power(-y2, dx1) * np.power(-y1, dx2)
223        if 'nearest' in method :
224            result = xr.where ( dx2>=dx1, y1, y2)
225
226        return result
227
228    for k in np.arange (nk_ou) :
229        result = __interp  (x[{pdim:k}], xp, yp)
230
231        # Put result in the final array
232        ou_tab [{pdim:k}] = result
233
234    return ou_tab.squeeze()
235
236def fixed_lon (lon, center_lon=0.0) :
237    '''Returns corrected longitudes for nicer plots
238
239    lon        : longitudes of the grid. At least 1D.
240    center_lon : center longitude. Default=0.
241
242    Designed by Phil Pelson. See https://gist.github.com/pelson/79cf31ef324774c97ae7
243    '''
244    mmath = __mmath__ (lon)
245
246    zfixed_lon = lon.copy ()
247
248    zfixed_lon = mmath.where (zfixed_lon > center_lon+180., zfixed_lon-360.0, fixed_lon)
249    zfixed_lon = mmath.where (zfixed_lon < center_lon-180., zfixed_lon+360.0, fixed_lon)
250
251    start = np.argmax (np.abs (np.diff (zfixed_lon, axis=-1)) > 180., axis=-1)
252    zfixed_lon [start+1:] += 360.
253
254    return zfixed_lon
255
256def nord2sud (p2d) :
257    '''Swap north to south a 2D field
258    '''
259    pout = p2d [..., -1::-1, : ]
260
261    return pout
262
263def point2geo (p1d, verbose=False, lon=False, lat=False, jpi=0, jpj=0, share_pole=False) :
264    '''
265    From 1D (..., points_physiques) (restart type) to 2D (..., lat, lon)
266    '''
267    math = __mmath__ (p1d)
268
269    # Get the horizontal dimension
270    jpn = p1d.shape[-1]
271    # Get other dimension
272    form1 = list(p1d.shape [0:-1]) 
273
274    if jpi==0 and jpj>0 : 
275        jpi = (jpn-2)//(jpj-2)
276    if jpi>0  and jpj == 0 :
277        jpj = (jpn-2)//jpi +2
278
279    def c_jpn (jpj, jpi) : return jpi*(jpj-2) + 2
280           
281    if jpi==0 and jpi==0 :
282        if jpn == c_jpn( 95, 96) :
283            jpj, jpi =  95,  96
284        if jpn == c_jpn (96, 96) :
285            jpj, jpi =  96,  96
286        if jpn == c_jpn (144, 144) :
287            jpj, jpi = 144, 144
288        if jpn == c_jpn (143, 144) :
289            jpj, jpi = 143, 144
290       
291    form2 = form1 + [jpj, jpi,]
292    form_shape = form1 + [jpj-2, jpi]
293    p2d = np.empty ( form2 )
294   
295    p2d [..., 1:jpj-1, :] = np.reshape (np.float64 (p1d [..., 1:jpn-1]), form_shape )
296    if share_pole :
297        for ji in np.arange ( p2d.shape[-1] ) :
298            p2d [..., 0      , ji]  = p1d[...,    0 ] / float(jpi)
299            p2d [..., jpj-1  , ji]  = p1d[..., jpn-1] / float(jpi)
300    else : 
301        for ji in np.arange ( p2d.shape[-1] ) :
302            p2d [..., 0      , ji]  = p1d[...,    0 ]
303            p2d [..., jpj-1  , ji]  = p1d[..., jpn-1]
304
305    if math == xr :
306        p2d = xr.DataArray (p2d)
307        p2d.attrs.update ( p1d.attrs )
308        for idim in np.arange ( len(p1d.shape [0:-1]) ):
309            dim = p1d.dims[idim]
310            p2d = p2d.rename        ( {p2d.dims[idim]:p1d.dims[idim]}  )
311            p2d = p2d.assign_coords ( {p2d.dims[idim]:p1d.coords[dim]} )
312        p2d = p2d.rename ( {p2d.dims[-1]:'x', p2d.dims[-2]:'y'} )
313        p2d['x'].attrs.update ( {'axis':'X'} )
314        p2d['y'].attrs.update ( {'axis':'Y'} )
315
316    if type(lon) == bool :
317        if lon :
318            lon = np.linspace ( -180, 180, jpi, endpoint=False)
319            #if math == xr :
320            #    lon = xr.DataArray ( lon, dims=('lon',), coords=(lon,) )
321    if type(lat) == bool :
322        if lat :
323            lat = np.linspace ( 90, -90, jpj, endpoint=True)       
324            #if math == xr :
325            #    lat = xr.DataArray ( lat, dims=('lat',), coords=(lat,) )
326    if type(lon) != bool and type(lat) != bool :
327        if  __mmath__ (lat) == xr :
328            lat_name = lat.name
329        else :
330            lat_name = 'lat'
331        if __mmath__ (lon) == xr : 
332            lon_name = lon.name
333        else :
334            lon_name = 'lon'
335        p2d = p2d.rename ( {p2d.dims[-1]:lon_name, p2d.dims[-2]:lat_name} )
336        p2d = p2d.assign_coords ( {lon_name:lon, lat_name:lat} )
337        p2d[lon_name].attrs.update ( { 'units':'degrees_east' , 'long_name':'Longitude', 'standard_name':'longitude', 'axis':'X' } )
338        p2d[lat_name].attrs.update ( { 'units':'degrees_north', 'long_name':'Latitude' , 'standard_name':'latitude' , 'axis':'Y' }  )
339
340    return p2d
341
342def geo2point ( p2d, cumul_poles=False, dim1d='points_physiques' ) :
343    '''
344    From 2D (..., lat, lon) to 1D (..., points_phyiques)
345    '''
346    math = __mmath__ (p2d)
347    #
348    # Get the horizontal dimensions
349    (jpj, jpi) = p2d.shape[-2:]
350    jpn = jpi*(jpj-2) + 2
351
352    # Get other dimensions
353    form1 = list(p2d.shape [0:-2]) 
354    form2 = form1 + [jpn,]
355       
356    p1d = np.empty ( form2 )
357    form_shape = form1 + [jpn-2,]
358   
359    if math == xr :
360            p1d[..., 1:-1] = np.reshape ( np.float64 (p2d[..., 1:-1, :].values).ravel(), form_shape )
361    else :
362            p1d[..., 1:-1] = np.reshape ( np.float64 (p2d[..., 1:-1, :]       ).ravel(), form_shape )
363    p1d[...,  0  ] = p2d[...,  0, 0]
364    p1d[..., -1  ] = p2d[..., -1, 0]
365
366    if math == xr :
367        p1d = xr.DataArray (p1d)
368        p1d.attrs.update ( p2d.attrs )
369        if len(p2d.shape [0:-2]) > 0 : 
370            for idim in np.arange ( len(p2d.shape [0:-2]) ):
371                dim=p2d.dims[idim]
372                p1d = p1d.rename        ( {p1d.dims[idim]:p2d.dims[idim]}   )
373                p1d = p1d.assign_coords ( {p1d.dims[idim] :p2d.coords[dim].values} )
374        p1d = p1d.rename ( {p1d.dims[-1]:dim1d} )
375
376    if cumul_poles :
377        p1d[...,  0] = np.sum ( p2d[...,  0, :] )
378        p1d[..., -1] = np.sum ( p2d[..., -1, :] )
379
380    return p1d
381
382def geo2en (pxx, pyy, pzz, glam, gphi) :
383    '''Change vector from geocentric to east/north
384
385    Inputs :
386        pxx, pyy, pzz : components on the geocentric system
387        glam, gphi : longitude and latitude of the points
388    '''
389
390    gsinlon = np.sin (RAD * glam)
391    gcoslon = np.cos (RAD * glam)
392    gsinlat = np.sin (RAD * gphi)
393    gcoslat = np.cos (RAD * gphi)
394
395    pte = - pxx * gsinlon            + pyy * gcoslon
396    ptn = - pxx * gcoslon * gsinlat  - pyy * gsinlon * gsinlat + pzz * gcoslat
397
398    return pte, ptn
399
400def en2geo (pte, ptn, glam, gphi) :
401    '''Change vector from east/north to geocentric
402
403    Inputs :
404        pte, ptn : eastward/northward components
405        glam, gphi : longitude and latitude of the points
406    '''
407
408    gsinlon = np.sin (RAD * glam)
409    gcoslon = np.cos (RAD * glam)
410    gsinlat = np.sin (RAD * gphi)
411    gcoslat = np.cos (RAD * gphi)
412
413    pxx = - pte * gsinlon - ptn * gcoslon * gsinlat
414    pyy =   pte * gcoslon - ptn * gsinlon * gsinlat
415    pzz =   ptn * gcoslat
416
417    return pxx, pyy, pzz
418
419## ===========================================================================
420##
421##                               That's all folk's !!!
422##
423## ===========================================================================
Note: See TracBrowser for help on using the repository browser.