acgc.geoschem API documentation

# acgc.geoschem

Tools for indexing, remapping, interpolating GEOS-Chem grids

View Source

      1#!/usr/bin/env python3
      2# -*- coding: utf-8 -*-
      3"""Tools for indexing, remapping, interpolating GEOS-Chem grids
      4"""
      5
      6import numpy as np
      7# from numba import jit
      8from scipy.interpolate import interp1d
      9
     10# @jit(nopython=True)
     11def mapweight1d( edge1, edge2, edgelump=True ):
     12    '''Calculate mapping weights from source grid with edge1 to destination grid with edge2
     13
     14    Parameters
     15    ----------
     16    edge1 : array (n,)
     17        source grid edge coordinates
     18    edge2 : array (m,)
     19        destination grid edge coordinates
     20    edgelump : bool, default=True)
     21        edgelump only affects situations where the source grid covers a larger domain
     22        than the destination grid (i.e. min(edge1)<min(edge2) or max(edge1)>max(edge2) ).
     23        If True, enforce total mass conservation in mapping weights by lumping excess mass
     24        in source grid into the corresponding end of destination grid.
     25        If False, the mass is conserved in the overlapping domains, but some total mass
     26        can be lost due to difference in overall domain.
     27
     28    Returns
     29    -------
     30    weight : array (m,n)
     31        mapping weights from source to destination grid: new = np.dot(weight,old)
     32    '''
     33
     34    reverse = False
     35
     36    # Check if the input pedge1 are inceasing
     37    if np.all(np.diff(edge1) > 0):
     38        # All are increasing, do nothing
     39        pass
     40    elif np.all(np.diff(edge1) < 0):
     41        raise ValueError( 'decreasing edge1 values not implemented')
     42        # All are decreasing, reverse
     43        edge1 = np.flip( edge1 )
     44    else:
     45        raise ValueError( 'Input edge1 must be sorted ')
     46
     47    # Check if the output edges are increasing
     48    if np.all(np.diff(edge2) > 0):
     49        # All are increasing, do nothing
     50        pass
     51    elif np.all(np.diff(edge2) < 0):
     52        raise ValueError( 'decreasing edge2 values not implemented')
     53        # All are decreasing, reverse
     54        edge2 = np.flip( edge2 )
     55        reverse = True
     56    else:
     57        raise ValueError( 'Input edge2 must be sorted ')
     58
     59    # Interpolate pedge2 into pedge1
     60    idx21 = np.interp( edge2, edge1, np.arange(len(edge1)) )
     61
     62    # If the pedge1 max or min are outside pedge2, then make sure the
     63    # edges go into the end levels
     64    if edgelump:
     65        idx21[0]  = 0
     66        idx21[-1] = len(edge1)-1
     67
     68    # Empty array for result
     69    weight = np.zeros( (len(edge2)-1, len(edge1)-1) )
     70
     71    for iNew in range(len(edge2)-1):
     72        for iOld in range( np.floor(idx21[iNew]).astype(int), np.ceil(idx21[iNew+1]).astype(int) ):
     73
     74            # Concise version
     75            weight[iNew,iOld] = ( np.minimum(idx21[iNew+1],iOld+1) - np.maximum(idx21[iNew],iOld) )
     76
     77    return weight
     78
     79def regrid2d( array1, xe1, ye1, xe2, ye2, edgelump=True ):
     80    '''Mass-conserving remapping of data between regular x-y or lat-lon grids
     81
     82    Parameters
     83    ----------
     84    array1 : array (n,m)
     85        input data. Quantities should be extensive (e.g. mass, not mixing ratio).
     86    xe1 : array (n+1,)
     87        source grid edges in x dimension, e.g. longitude
     88    ye1 : array (m+1,)
     89        source grid edges in y dimension, e.g. latitude
     90    xe2 : array (p+1,)
     91        destination grid edges in x dimension, e.g. longitude
     92    ye2 : array (q+1,)
     93        destination grid edges in y dimension, e.g. latitude
     94    edgelump : bool, default=True
     95        edgelump only affects situations where the source grid covers a larger domain
     96        than the destination grid (i.e. min(edge1)<min(edge2) or max(edge1)>max(edge2) ).
     97        If True, enforce total mass conservation in mapping weights by lumping excess mass
     98        in source grid into the corresponding end of destination grid.
     99        If False, the mass is conserved in the overlapping domains, but some total mass
    100        can be lost due to difference in overall domain.
    101
    102    Returns
    103    -------
    104    outarray : array (p,q)
    105        input data remapped to the destination grid
    106    '''
    107
    108    # Sizes
    109    insize = np.shape(array1)
    110    nxe1 = np.size(xe1)
    111    nye1 = np.size(ye1)
    112    nxe2 = np.size(xe2)
    113    nye2 = np.size(ye2)
    114
    115    # Ensure arrays conform
    116    if len(insize) != 2 :
    117        raise TypeError( 'array1 must be 2D')
    118    if insize[1] != nxe1-1:
    119        raise ValueError( 'xe1 must match size of array1' )
    120    if insize[0] != nye1-1:
    121        raise ValueError( 'ye1 must match size of array1' )
    122
    123    # Regridding weight matrices
    124    weightY = mapweight1d(ye1, ye2, edgelump)
    125    weightX = mapweight1d(xe1, xe2, edgelump)
    126
    127    # Size of temporary and output arrays
    128    tmpsize = [nye2-1,nxe1-1]
    129    outsize = [nye2-1,nxe2-1]
    130
    131    # Temporary array
    132    a1b = np.empty( tmpsize )
    133    a1b[:] = np.nan
    134
    135    # Output array
    136    outarray = np.empty(outsize)
    137    outarray[:] = np.nan
    138
    139    # Regrid latitude
    140    for i in range(np.size(xe1)-1):
    141        a1b[:,i] = np.dot( weightY, array1[:,i].flatten() )
    142
    143    # Regrid longitude
    144    for i in range(outsize[0]):
    145        outarray[i,:] = np.dot( weightX, a1b[i,:].flatten() )
    146
    147    return outarray
    148
    149# @jit(nopython=True)
    150def regrid_plevels( array1, pedge1, pedge2, intensive=False, edgelump=True ):
    151    ''' Mass-conserving regridding of data on pressure levels
    152    Data from the input array will be remapped to the output pressure levels while
    153    conserving total mass.
    154
    155    Parameters
    156    ----------
    157    array1 : array, (p,)
    158        Array of data that will be remapped.
    159        Quantities should be extensive (e.g. mass, not mixing ratio) unless the
    160        intensive keyword is used.
    161    pedge1 : array (p+1,)
    162        Pressure edges of the input array1. The length of pedge1 should be one greater than array1.
    163        The element array1[i] is bounded by edges pedge1[i], pedge1[i+1]
    164    pedge2 : array (q+1,)
    165        Pressure edges of the desired output array.
    166    intensive : bool, default=False
    167        If True, then array1 will be treated as an intensive quantity (e.g mole fraction,
    168        mixing ratio) during regridding. Each element array1[i] will be multiplied by
    169        (pedge1[i]-pedge[i+1]), which is proportional to airmass if pedge1 is pressure,
    170        before regridding. The output array2, will similarly be dividied by (pedge2[i]-pedge[i+1])
    171        so that the output is also intensive.
    172        By default, intensive=False is used and the input array is assumed to be extensive.
    173    edgelump : bool, default=True
    174        If the max or min of pedge1 extend beyond the max or min of pedge2 and edgelump=True,
    175        the mass will be placed in in the first or last grid level of the output array2,
    176        ensuring that the mass of array2 is the same as array1.
    177        If edgelump=False, then the mass of array2 can be less than array1.
    178
    179    Returns
    180    array2 : array (q,)
    181        Array of remapped data. The length of array2 will be one less than pedge2.
    182    '''
    183
    184    if len(pedge1) != len(array1)+1:
    185        raise ValueError("pedge1 must have size len(array1)+1")
    186
    187    reverse = False
    188
    189    # Check if the input pedge1 are inceasing
    190    if np.all(np.diff(pedge1) > 0):
    191        # All are increasing, do nothing
    192        pass
    193    elif np.all(np.diff(pedge1) < 0):
    194        # All are decreasing, reverse
    195        array1 = np.flip( array1 )
    196        pedge1 = np.flip( pedge1 )
    197    else:
    198        raise ValueError( 'Input pedge1 must be sorted ')
    199
    200    # Check if the output edges are increasing
    201    if np.all(np.diff(pedge2) > 0):
    202        # All are increasing, do nothing
    203        pass
    204    elif np.all(np.diff(pedge2) < 0):
    205        # All are decreasing, reverse
    206        pedge2 = np.flip( pedge2 )
    207        reverse = True
    208    else:
    209        raise ValueError( 'Input pedge2 must be sorted ')
    210
    211    # If input is an intensive quantity, then multiply by the pressure grid spacing
    212    if intensive:
    213        array1 = array1 * ( pedge1[1:] - pedge1[:-1] )
    214
    215    # Interpolate pedge2 into pedge1
    216    idx21 = np.interp( pedge2, pedge1, np.arange(len(pedge1)) )
    217
    218    # If the pedge1 max or min are outside pedge2, then make sure the
    219    # edges go into the end levels
    220    if edgelump:
    221        idx21[0]  = 0
    222        idx21[-1] = len(pedge1)-1
    223
    224    # Empty array for result
    225    array2 = np.zeros( len(pedge2)-1 )
    226
    227    for i in range(len(array2)):
    228        for j in range( np.floor(idx21[i]).astype(int), np.ceil(idx21[i+1]).astype(int) ):
    229
    230            # Concise version
    231            array2[i] += array1[j] * ( np.minimum(idx21[i+1],j+1) - np.maximum(idx21[i],j) )
    232
    233            # Lengthy version
    234            #if ( j >= idx21[i] and j+1 <= idx21[i+1] ):
    235            #    array2[i] += array1[j]
    236            #    #print( 'input level {:d} inside output level {:d}'.format(j,i))
    237            #elif ( j < idx21[i] and j+1 > idx21[i+1] ):
    238            #    array2[i] += array1[j] * (idx21[i+1] - idx21[i])
    239            #    #print( 'output level inside input level')
    240            #elif ( j < idx21[i] ):
    241            #    array2[i] +=  array1[j] * (1 - np.mod(idx21[i],  1))
    242            #    #print( 'input level {:d} at bottom of output level {:d}'.format(j,i))
    243            #elif ( j+1 > idx21[i+1] ):
    244            #    array2[i] += array1[j] * np.mod( idx21[i+1], 1)
    245            #    #print( 'input level {:d} at top of output level {:d}'.format(j,i))
    246            #else:
    247            #    raise NotImplemented
    248
    249    # Convert back to an intensive quantity, if necessary
    250    if intensive:
    251        array2 = array2 / ( pedge2[1:] - pedge2[:-1] )
    252
    253    # Reverse the output array, if necessary
    254    if reverse:
    255        array2 = np.flip( array2 )
    256
    257    return array2
    258
    259def _set_center_edge( center, edge, name ):
    260    '''Resolve any conflict between center and edge keywords
    261
    262    Arguments
    263    ---------
    264    center : bool or None
    265    edge   : bool or None
    266    name : str
    267        name of the calling function, only used in error messages
    268
    269    Returns
    270    -------
    271    center, edge : bool
    272        consistent values: one True and the other False
    273    '''
    274
    275    if (center is None) and (edge is None):
    276
    277        # If both center and edge are None, default to center
    278        center = True
    279        edge   = False
    280
    281    elif center is None:
    282
    283        # Use edge, if that value is boolean
    284        if isinstance(edge, bool):
    285            center = not edge
    286        else:
    287            raise TypeError( name + ': edge must be boolean or None')
    288
    289    elif edge is None:
    290
    291        # Use center, if that value is boolean
    292        if isinstance(center, bool):
    293            edge = not center
    294        else:
    295            raise TypeError( name + ': center must be boolean or None')
    296
    297    elif (isinstance(center, bool) and isinstance(edge, bool)):
    298
    299        # If both are True or both are False, then raise error
    300        if ( ((center is True)  and (edge is True)) or
    301             ((center is False) and (edge is False)) ):
    302            # If both center and edge are specified, then raise error
    303            raise ValueError( name + ': Choose either edge=True or center=True' )
    304
    305        # At this point, either center=True or edge=True, but not both
    306
    307    else:
    308
    309        # Values aren't boolean
    310        raise TypeError( name + ': center and edge must be boolean or None')
    311
    312    return center, edge
    313
    314def lon2i(lon,res=4):
    315    '''Find I grid index for longitude
    316
    317    Parameters
    318    ----------
    319    lon : float
    320        longitude in degrees
    321    res : float, default=4
    322        resolution of grid. See `get_lon` for allowed values
    323
    324    Result
    325    ------
    326    index : int
    327        grid index location closest to longitude
    328    '''
    329    # Put longitude in expected GC the range [-180,180]
    330    alon = np.mod( np.asarray(lon) + 180, 360 ) - 180
    331    longrid = get_lon(res=res,center=True)
    332    nlon = len(longrid)
    333    f = interp1d(np.append(longrid,longrid[0]+360),
    334                np.append(np.arange(nlon),0),
    335                kind='nearest',
    336                fill_value='extrapolate')
    337    return f(alon).astype(np.int32)
    338
    339def lat2j(lat,res=4):
    340    '''Find J grid index for latitude
    341
    342    Parameters
    343    ----------
    344    lat : float
    345        latitude in degrees
    346    res : float, default=4
    347        resolution of grid. See `get_lon` for allowed values
    348
    349    Result
    350    ------
    351    index : int
    352        grid index location closest to latitude
    353    '''
    354    alat = np.asarray(lat)
    355    latgrid = get_lat(res=res,center=True)
    356    nlat = len(latgrid)
    357    if (np.any(alat<-90) or np.any(alat>90)):
    358        raise ValueError('lat must be in the range [-90,90]')
    359    f = interp1d(latgrid,
    360                np.arange(nlat),
    361                kind='nearest',
    362                fill_value='extrapolate')
    363    return f(alat).astype(np.int32)
    364
    365def ll2ij(lon,lat,res=4):
    366    '''Find (I,J) grid index for longitude, latitude
    367
    368
    369    Parameters
    370    ----------
    371    lon, lat : float
    372        longitude and latitude in degrees
    373    res : float, default=4
    374        resolution of grid. See `get_lon` for allowed values
    375
    376    Result
    377    ------
    378    i, j : int
    379        grid index location closest to longitude, latitude
    380    '''
    381    alon = np.asarray(lon)
    382    alat = np.asarray(lat)
    383    nlon = alon.size
    384    nlat = alat.size
    385    if (nlon==1 and nlat>1):
    386        alon = np.ones_like(alat) * alon
    387    elif (nlat==1 and nlon>1):
    388        alat = np.ones_like(alon) * alat
    389    elif nlon != nlat:
    390        raise ValueError('lon and lat must have the same length or be conformable')
    391    return lon2i(alon,res=res), lat2j(alat,res=res)
    392
    393def get_lon(res=4, center=None, edge=None):
    394    '''Get array of longitudes for grid centers or edges.
    395
    396    GMAO grids assumed
    397    - res = 0.5 for 0.5x0.625
    398    - res = 2 for 2x2.5
    399    - res = 4 for 4x5
    400
    401    Parameters
    402    ----------
    403    res : float, default=4
    404        latitude resolution of grid
    405    center, edge : bool or None, default=None
    406        determines whether returned values are at grid center (center=True) or grid edge (edge=True)
    407        If both center and edge are None, then center is used
    408
    409    Returns
    410    -------
    411    lon : array (n,)
    412        longitude in degrees at grid center or edges
    413    '''
    414
    415    # Resolve any conflict between center and edge
    416    center, edge = _set_center_edge( center, edge, 'get_lon' )
    417
    418    # Longitude edges
    419    if res==2:
    420        dx=2.5
    421    elif res==4:
    422        dx=5
    423    elif res==0.5:
    424        dx=0.625
    425    else:
    426        raise NotImplementedError( 'Resolution '+str(res)+' is not defined' )
    427
    428    # Longitude edges
    429    lonedge = np.arange(-180 - dx/2, 180, dx)
    430
    431    # Get grid centers, if needed
    432    if center:
    433        lon = ( lonedge[0:-1] + lonedge[1:] ) / 2
    434    else:
    435        lon = lonedge
    436
    437    return lon
    438
    439def get_lat(res=4, center=None, edge=None):
    440    '''Get array of latitudes at grid center or edges
    441
    442    GMAO grids assumed
    443    - res = 0.5 for 0.5x0.625
    444    - res = 2 for 2x2.5
    445    - res = 4 for 4x5
    446
    447    Parameters
    448    ----------
    449    res : float, default=4
    450        latitude resolution of grid
    451    center, edge : bool or None, default=None
    452        determines whether returned values are at grid center (center=True) or grid edge (edge=True)
    453        If both center and edge are None, then center is used
    454
    455    Returns
    456    -------
    457    lat : array (n,)
    458        latitude in degrees at grid center or edges
    459    '''
    460
    461    # Resolve any conflict between center and edge
    462    center, edge = _set_center_edge( center, edge, 'get_lat' )
    463
    464    # Latitude edges
    465    if res==2:
    466        dy=2
    467    elif res==4:
    468        dy=4
    469    elif res==0.5:
    470        dy=0.5
    471    else:
    472        raise NotImplementedError( 'Resolution '+str(res)+' is not defined' )
    473
    474    # Latitude edges
    475    latedge = np.hstack([-90, np.arange(-90+dy/2,90,dy), 90] )
    476
    477    # Get grid centers, if needed
    478    if center:
    479        lat = ( latedge[0:-1] + latedge[1:] ) / 2
    480    else:
    481        lat = latedge
    482
    483    return lat
    484
    485def get_lev_p(p_surf=1013.25, nlev=47, edge=False ):
    486    '''Get array of pressures at grid centers or edges
    487
    488    Parameters
    489    ----------
    490    p_surf : float
    491        surface pressure, hPa
    492    nlev : {47,72}
    493        number of vertical levels (47 or 72)
    494    edge : bool
    495        specify edges (True) or centers (False, default)
    496
    497    Returns
    498    -------
    499    p : array of float
    500        pressures for each vertical level'''
    501
    502    # Get A,B values for vertical coordinate
    503    A, B = get_hybrid_ab(nlev=nlev, edge=edge )
    504
    505    # Pressure, hPa
    506    P = A + B * p_surf
    507
    508    return P
    509
    510# @jit(nopython=True)
    511def get_lev_z(TK=[273], p_surf=1013.25, p_edge=[0], Q=0, z_surf=0, edge=False ):
    512    '''Calculate the geopotential heights of level edges or centers
    513
    514    Calculations use hypsometric equation, assuing hydrostatic balance
    515
    516    Parameters
    517    ----------
    518    TK : array of float
    519        Temperature for each grid level, K
    520    p_surf : float
    521        Surface pressure, hPa
    522    p_edge : array of float, optional
    523        Pressure at each grid level edge, hPa
    524        Will be calculated if length differs from TK
    525    Q : array of float
    526        Specific humidity for each grid level, kg/kg
    527    z_surf : float
    528        Surface altitude, m
    529    edge : bool
    530        specify edges (True) or centers (False, default) for output
    531
    532    Returns
    533    -------
    534    z : array of float
    535        Geopotential height for each grid edge or center, m
    536    '''
    537
    538    funname='geoschem: get_lev_z: '
    539
    540    # Acceleration due to gravity, m/s2
    541    g0 = 9.81
    542
    543    if len(TK)<=1:
    544        print(funname,'TK must be an array')
    545        return -1
    546
    547    # Number of levels in the temperature profile
    548    nlev = TK.size
    549
    550    # Calculate the pressure edges if they aren't provided in input
    551    if len(p_edge) != nlev+1:
    552        # Get pressure at level edges, hPa
    553        p_edge = get_lev_p(p_surf=p_surf, nlev=nlev, edge=True )
    554
    555    # Initialize array
    556    z_edge = np.zeros_like(p_edge)
    557    z_edge[0] = z_surf
    558
    559    # Virtual temperature, K
    560    Tv = TK * ( 1 + 0.61 * Q )
    561
    562    # Calculate altitudes with hypsometric equation, m
    563    for L in range(1, p_edge.size):
    564        z_edge[L] = z_edge[L-1] + 287 * Tv[L-1] / g0 * np.log( p_edge[L-1] / p_edge[L] )
    565
    566    # Choose grid edge or center altitudes, m
    567    if edge is True:
    568        z = z_edge
    569    else:
    570        z = (z_edge[0:-1] + z_edge[1:]) / 2
    571
    572    return z
    573
    574def get_hybrid_ab(nlev=47, edge=False, center=False):
    575    '''Get A and B values for the GEOS-5/MERRA/MERRA2 hybrid sigma-pressure vertical coordinate.
    576
    577    The pressure at level edge L is P(L) = A(L) + B(L) * PS, where PS is surface pressure in hPa
    578
    579    Parameters
    580    ----------
    581    nlev : {47,72}
    582        number of levels in the desired grid (47 or 72)
    583    center, edge : bool
    584        specify if values for level center or edges are desired
    585
    586    Returns
    587    -------
    588    A, B : array (n,)
    589    '''
    590
    591    # Resolve any conflict between center and edge requests
    592    center, edge = _set_center_edge( center, edge, 'get_hybrid_ab')
    593
    594    if nlev == 47:
    595
    596        # A [hPa] for 47 levels (48 edges)
    597        A = np.array([ 0.000000E+00, 4.804826E-02, 6.593752E+00, 1.313480E+01,
    598               1.961311E+01, 2.609201E+01, 3.257081E+01, 3.898201E+01,
    599               4.533901E+01, 5.169611E+01, 5.805321E+01, 6.436264E+01,
    600               7.062198E+01, 7.883422E+01, 8.909992E+01, 9.936521E+01,
    601               1.091817E+02, 1.189586E+02, 1.286959E+02, 1.429100E+02,
    602               1.562600E+02, 1.696090E+02, 1.816190E+02, 1.930970E+02,
    603               2.032590E+02, 2.121500E+02, 2.187760E+02, 2.238980E+02,
    604               2.243630E+02, 2.168650E+02, 2.011920E+02, 1.769300E+02,
    605               1.503930E+02, 1.278370E+02, 1.086630E+02, 9.236572E+01,
    606               7.851231E+01, 5.638791E+01, 4.017541E+01, 2.836781E+01,
    607               1.979160E+01, 9.292942E+00, 4.076571E+00, 1.650790E+00,
    608               6.167791E-01, 2.113490E-01, 6.600001E-02, 1.000000E-02 ] )
    609
    610        # Bp [unitless] for 47 levels (48 edges)
    611        B = np.array([ 1.000000E+00, 9.849520E-01, 9.634060E-01, 9.418650E-01,
    612               9.203870E-01, 8.989080E-01, 8.774290E-01, 8.560180E-01,
    613               8.346609E-01, 8.133039E-01, 7.919469E-01, 7.706375E-01,
    614               7.493782E-01, 7.211660E-01, 6.858999E-01, 6.506349E-01,
    615               6.158184E-01, 5.810415E-01, 5.463042E-01, 4.945902E-01,
    616               4.437402E-01, 3.928911E-01, 3.433811E-01, 2.944031E-01,
    617               2.467411E-01, 2.003501E-01, 1.562241E-01, 1.136021E-01,
    618               6.372006E-02, 2.801004E-02, 6.960025E-03, 8.175413E-09,
    619               0.000000E+00, 0.000000E+00, 0.000000E+00, 0.000000E+00,
    620               0.000000E+00, 0.000000E+00, 0.000000E+00, 0.000000E+00,
    621               0.000000E+00, 0.000000E+00, 0.000000E+00, 0.000000E+00,
    622               0.000000E+00, 0.000000E+00, 0.000000E+00, 0.000000E+00 ] )
    623
    624    elif nlev == 72:
    625
    626        # A [hPa] for 72 levels (73 edges)
    627        A = np.array([ 0.000000E+00, 4.804826E-02, 6.593752E+00, 1.313480E+01,
    628               1.961311E+01, 2.609201E+01, 3.257081E+01, 3.898201E+01,
    629               4.533901E+01, 5.169611E+01, 5.805321E+01, 6.436264E+01,
    630               7.062198E+01, 7.883422E+01, 8.909992E+01, 9.936521E+01,
    631               1.091817E+02, 1.189586E+02, 1.286959E+02, 1.429100E+02,
    632               1.562600E+02, 1.696090E+02, 1.816190E+02, 1.930970E+02,
    633               2.032590E+02, 2.121500E+02, 2.187760E+02, 2.238980E+02,
    634               2.243630E+02, 2.168650E+02, 2.011920E+02, 1.769300E+02,
    635               1.503930E+02, 1.278370E+02, 1.086630E+02, 9.236572E+01,
    636               7.851231E+01, 6.660341E+01, 5.638791E+01, 4.764391E+01,
    637               4.017541E+01, 3.381001E+01, 2.836781E+01, 2.373041E+01,
    638               1.979160E+01, 1.645710E+01, 1.364340E+01, 1.127690E+01,
    639               9.292942E+00, 7.619842E+00, 6.216801E+00, 5.046801E+00,
    640               4.076571E+00, 3.276431E+00, 2.620211E+00, 2.084970E+00,
    641               1.650790E+00, 1.300510E+00, 1.019440E+00, 7.951341E-01,
    642               6.167791E-01, 4.758061E-01, 3.650411E-01, 2.785261E-01,
    643               2.113490E-01, 1.594950E-01, 1.197030E-01, 8.934502E-02,
    644               6.600001E-02, 4.758501E-02, 3.270000E-02, 2.000000E-02,
    645               1.000000E-02 ] )
    646
    647        # B [unitless] for 72 levels (73 edges)
    648        B = np.array([ 1.000000E+00, 9.849520E-01, 9.634060E-01, 9.418650E-01,
    649               9.203870E-01, 8.989080E-01, 8.774290E-01, 8.560180E-01,
    650               8.346609E-01, 8.133039E-01, 7.919469E-01, 7.706375E-01,
    651               7.493782E-01, 7.211660E-01, 6.858999E-01, 6.506349E-01,
    652               6.158184E-01, 5.810415E-01, 5.463042E-01, 4.945902E-01,
    653               4.437402E-01, 3.928911E-01, 3.433811E-01, 2.944031E-01,
    654               2.467411E-01, 2.003501E-01, 1.562241E-01, 1.136021E-01,
    655               6.372006E-02, 2.801004E-02, 6.960025E-03, 8.175413E-09,
    656               0.000000E+00, 0.000000E+00, 0.000000E+00, 0.000000E+00,
    657               0.000000E+00, 0.000000E+00, 0.000000E+00, 0.000000E+00,
    658               0.000000E+00, 0.000000E+00, 0.000000E+00, 0.000000E+00,
    659               0.000000E+00, 0.000000E+00, 0.000000E+00, 0.000000E+00,
    660               0.000000E+00, 0.000000E+00, 0.000000E+00, 0.000000E+00,
    661               0.000000E+00, 0.000000E+00, 0.000000E+00, 0.000000E+00,
    662               0.000000E+00, 0.000000E+00, 0.000000E+00, 0.000000E+00,
    663               0.000000E+00, 0.000000E+00, 0.000000E+00, 0.000000E+00,
    664               0.000000E+00, 0.000000E+00, 0.000000E+00, 0.000000E+00,
    665               0.000000E+00, 0.000000E+00, 0.000000E+00, 0.000000E+00,
    666               0.000000E+00 ])
    667
    668    else:
    669        raise ValueError( f'get_hybrid_ab: A, B not defined for nlev={nlev}' )
    670
    671    # Calculate values at level center
    672    if center is True:
    673
    674        A = (A[0:-1] + A[1:]) / 2
    675        B = (B[0:-1] + B[1:]) / 2
    676
    677    return A, B

def mapweight1d(edge1, edge2, edgelump=True): View Source

    12def mapweight1d( edge1, edge2, edgelump=True ):
    13    '''Calculate mapping weights from source grid with edge1 to destination grid with edge2
    14
    15    Parameters
    16    ----------
    17    edge1 : array (n,)
    18        source grid edge coordinates
    19    edge2 : array (m,)
    20        destination grid edge coordinates
    21    edgelump : bool, default=True)
    22        edgelump only affects situations where the source grid covers a larger domain
    23        than the destination grid (i.e. min(edge1)<min(edge2) or max(edge1)>max(edge2) ).
    24        If True, enforce total mass conservation in mapping weights by lumping excess mass
    25        in source grid into the corresponding end of destination grid.
    26        If False, the mass is conserved in the overlapping domains, but some total mass
    27        can be lost due to difference in overall domain.
    28
    29    Returns
    30    -------
    31    weight : array (m,n)
    32        mapping weights from source to destination grid: new = np.dot(weight,old)
    33    '''
    34
    35    reverse = False
    36
    37    # Check if the input pedge1 are inceasing
    38    if np.all(np.diff(edge1) > 0):
    39        # All are increasing, do nothing
    40        pass
    41    elif np.all(np.diff(edge1) < 0):
    42        raise ValueError( 'decreasing edge1 values not implemented')
    43        # All are decreasing, reverse
    44        edge1 = np.flip( edge1 )
    45    else:
    46        raise ValueError( 'Input edge1 must be sorted ')
    47
    48    # Check if the output edges are increasing
    49    if np.all(np.diff(edge2) > 0):
    50        # All are increasing, do nothing
    51        pass
    52    elif np.all(np.diff(edge2) < 0):
    53        raise ValueError( 'decreasing edge2 values not implemented')
    54        # All are decreasing, reverse
    55        edge2 = np.flip( edge2 )
    56        reverse = True
    57    else:
    58        raise ValueError( 'Input edge2 must be sorted ')
    59
    60    # Interpolate pedge2 into pedge1
    61    idx21 = np.interp( edge2, edge1, np.arange(len(edge1)) )
    62
    63    # If the pedge1 max or min are outside pedge2, then make sure the
    64    # edges go into the end levels
    65    if edgelump:
    66        idx21[0]  = 0
    67        idx21[-1] = len(edge1)-1
    68
    69    # Empty array for result
    70    weight = np.zeros( (len(edge2)-1, len(edge1)-1) )
    71
    72    for iNew in range(len(edge2)-1):
    73        for iOld in range( np.floor(idx21[iNew]).astype(int), np.ceil(idx21[iNew+1]).astype(int) ):
    74
    75            # Concise version
    76            weight[iNew,iOld] = ( np.minimum(idx21[iNew+1],iOld+1) - np.maximum(idx21[iNew],iOld) )
    77
    78    return weight

Calculate mapping weights from source grid with edge1 to destination grid with edge2

###### Parameters

- edge1 (array (n,)): source grid edge coordinates
- edge2 (array (m,)): destination grid edge coordinates
- edgelump (bool, default=True)): edgelump only affects situations where the source grid covers a larger domain than the destination grid (i.e. min(edge1)

###### Returns

- weight (array (m,n)): mapping weights from source to destination grid: new = np.dot(weight,old)

def regrid2d(array1, xe1, ye1, xe2, ye2, edgelump=True): View Source

     80def regrid2d( array1, xe1, ye1, xe2, ye2, edgelump=True ):
     81    '''Mass-conserving remapping of data between regular x-y or lat-lon grids
     82
     83    Parameters
     84    ----------
     85    array1 : array (n,m)
     86        input data. Quantities should be extensive (e.g. mass, not mixing ratio).
     87    xe1 : array (n+1,)
     88        source grid edges in x dimension, e.g. longitude
     89    ye1 : array (m+1,)
     90        source grid edges in y dimension, e.g. latitude
     91    xe2 : array (p+1,)
     92        destination grid edges in x dimension, e.g. longitude
     93    ye2 : array (q+1,)
     94        destination grid edges in y dimension, e.g. latitude
     95    edgelump : bool, default=True
     96        edgelump only affects situations where the source grid covers a larger domain
     97        than the destination grid (i.e. min(edge1)<min(edge2) or max(edge1)>max(edge2) ).
     98        If True, enforce total mass conservation in mapping weights by lumping excess mass
     99        in source grid into the corresponding end of destination grid.
    100        If False, the mass is conserved in the overlapping domains, but some total mass
    101        can be lost due to difference in overall domain.
    102
    103    Returns
    104    -------
    105    outarray : array (p,q)
    106        input data remapped to the destination grid
    107    '''
    108
    109    # Sizes
    110    insize = np.shape(array1)
    111    nxe1 = np.size(xe1)
    112    nye1 = np.size(ye1)
    113    nxe2 = np.size(xe2)
    114    nye2 = np.size(ye2)
    115
    116    # Ensure arrays conform
    117    if len(insize) != 2 :
    118        raise TypeError( 'array1 must be 2D')
    119    if insize[1] != nxe1-1:
    120        raise ValueError( 'xe1 must match size of array1' )
    121    if insize[0] != nye1-1:
    122        raise ValueError( 'ye1 must match size of array1' )
    123
    124    # Regridding weight matrices
    125    weightY = mapweight1d(ye1, ye2, edgelump)
    126    weightX = mapweight1d(xe1, xe2, edgelump)
    127
    128    # Size of temporary and output arrays
    129    tmpsize = [nye2-1,nxe1-1]
    130    outsize = [nye2-1,nxe2-1]
    131
    132    # Temporary array
    133    a1b = np.empty( tmpsize )
    134    a1b[:] = np.nan
    135
    136    # Output array
    137    outarray = np.empty(outsize)
    138    outarray[:] = np.nan
    139
    140    # Regrid latitude
    141    for i in range(np.size(xe1)-1):
    142        a1b[:,i] = np.dot( weightY, array1[:,i].flatten() )
    143
    144    # Regrid longitude
    145    for i in range(outsize[0]):
    146        outarray[i,:] = np.dot( weightX, a1b[i,:].flatten() )
    147
    148    return outarray

Mass-conserving remapping of data between regular x-y or lat-lon grids

###### Parameters

- array1 (array (n,m)): input data. Quantities should be extensive (e.g. mass, not mixing ratio).
- xe1 (array (n+1,)): source grid edges in x dimension, e.g. longitude
- ye1 (array (m+1,)): source grid edges in y dimension, e.g. latitude
- xe2 (array (p+1,)): destination grid edges in x dimension, e.g. longitude
- ye2 (array (q+1,)): destination grid edges in y dimension, e.g. latitude
- edgelump (bool, default=True): edgelump only affects situations where the source grid covers a larger domain than the destination grid (i.e. min(edge1)

###### Returns

- outarray (array (p,q)): input data remapped to the destination grid

def regrid_plevels(array1, pedge1, pedge2, intensive=False, edgelump=True): View Source

    151def regrid_plevels( array1, pedge1, pedge2, intensive=False, edgelump=True ):
    152    ''' Mass-conserving regridding of data on pressure levels
    153    Data from the input array will be remapped to the output pressure levels while
    154    conserving total mass.
    155
    156    Parameters
    157    ----------
    158    array1 : array, (p,)
    159        Array of data that will be remapped.
    160        Quantities should be extensive (e.g. mass, not mixing ratio) unless the
    161        intensive keyword is used.
    162    pedge1 : array (p+1,)
    163        Pressure edges of the input array1. The length of pedge1 should be one greater than array1.
    164        The element array1[i] is bounded by edges pedge1[i], pedge1[i+1]
    165    pedge2 : array (q+1,)
    166        Pressure edges of the desired output array.
    167    intensive : bool, default=False
    168        If True, then array1 will be treated as an intensive quantity (e.g mole fraction,
    169        mixing ratio) during regridding. Each element array1[i] will be multiplied by
    170        (pedge1[i]-pedge[i+1]), which is proportional to airmass if pedge1 is pressure,
    171        before regridding. The output array2, will similarly be dividied by (pedge2[i]-pedge[i+1])
    172        so that the output is also intensive.
    173        By default, intensive=False is used and the input array is assumed to be extensive.
    174    edgelump : bool, default=True
    175        If the max or min of pedge1 extend beyond the max or min of pedge2 and edgelump=True,
    176        the mass will be placed in in the first or last grid level of the output array2,
    177        ensuring that the mass of array2 is the same as array1.
    178        If edgelump=False, then the mass of array2 can be less than array1.
    179
    180    Returns
    181    array2 : array (q,)
    182        Array of remapped data. The length of array2 will be one less than pedge2.
    183    '''
    184
    185    if len(pedge1) != len(array1)+1:
    186        raise ValueError("pedge1 must have size len(array1)+1")
    187
    188    reverse = False
    189
    190    # Check if the input pedge1 are inceasing
    191    if np.all(np.diff(pedge1) > 0):
    192        # All are increasing, do nothing
    193        pass
    194    elif np.all(np.diff(pedge1) < 0):
    195        # All are decreasing, reverse
    196        array1 = np.flip( array1 )
    197        pedge1 = np.flip( pedge1 )
    198    else:
    199        raise ValueError( 'Input pedge1 must be sorted ')
    200
    201    # Check if the output edges are increasing
    202    if np.all(np.diff(pedge2) > 0):
    203        # All are increasing, do nothing
    204        pass
    205    elif np.all(np.diff(pedge2) < 0):
    206        # All are decreasing, reverse
    207        pedge2 = np.flip( pedge2 )
    208        reverse = True
    209    else:
    210        raise ValueError( 'Input pedge2 must be sorted ')
    211
    212    # If input is an intensive quantity, then multiply by the pressure grid spacing
    213    if intensive:
    214        array1 = array1 * ( pedge1[1:] - pedge1[:-1] )
    215
    216    # Interpolate pedge2 into pedge1
    217    idx21 = np.interp( pedge2, pedge1, np.arange(len(pedge1)) )
    218
    219    # If the pedge1 max or min are outside pedge2, then make sure the
    220    # edges go into the end levels
    221    if edgelump:
    222        idx21[0]  = 0
    223        idx21[-1] = len(pedge1)-1
    224
    225    # Empty array for result
    226    array2 = np.zeros( len(pedge2)-1 )
    227
    228    for i in range(len(array2)):
    229        for j in range( np.floor(idx21[i]).astype(int), np.ceil(idx21[i+1]).astype(int) ):
    230
    231            # Concise version
    232            array2[i] += array1[j] * ( np.minimum(idx21[i+1],j+1) - np.maximum(idx21[i],j) )
    233
    234            # Lengthy version
    235            #if ( j >= idx21[i] and j+1 <= idx21[i+1] ):
    236            #    array2[i] += array1[j]
    237            #    #print( 'input level {:d} inside output level {:d}'.format(j,i))
    238            #elif ( j < idx21[i] and j+1 > idx21[i+1] ):
    239            #    array2[i] += array1[j] * (idx21[i+1] - idx21[i])
    240            #    #print( 'output level inside input level')
    241            #elif ( j < idx21[i] ):
    242            #    array2[i] +=  array1[j] * (1 - np.mod(idx21[i],  1))
    243            #    #print( 'input level {:d} at bottom of output level {:d}'.format(j,i))
    244            #elif ( j+1 > idx21[i+1] ):
    245            #    array2[i] += array1[j] * np.mod( idx21[i+1], 1)
    246            #    #print( 'input level {:d} at top of output level {:d}'.format(j,i))
    247            #else:
    248            #    raise NotImplemented
    249
    250    # Convert back to an intensive quantity, if necessary
    251    if intensive:
    252        array2 = array2 / ( pedge2[1:] - pedge2[:-1] )
    253
    254    # Reverse the output array, if necessary
    255    if reverse:
    256        array2 = np.flip( array2 )
    257
    258    return array2

Mass-conserving regridding of data on pressure levels Data from the input array will be remapped to the output pressure levels while conserving total mass.

###### Parameters

- array1 (array, (p,)): Array of data that will be remapped. Quantities should be extensive (e.g. mass, not mixing ratio) unless the intensive keyword is used.
- pedge1 (array (p+1,)): Pressure edges of the input array1. The length of pedge1 should be one greater than array1. The element array1[i] is bounded by edges pedge1[i], pedge1[i+1]
- pedge2 (array (q+1,)): Pressure edges of the desired output array.
- intensive (bool, default=False): If True, then array1 will be treated as an intensive quantity (e.g mole fraction, mixing ratio) during regridding. Each element array1[i] will be multiplied by (pedge1[i]-pedge[i+1]), which is proportional to airmass if pedge1 is pressure, before regridding. The output array2, will similarly be dividied by (pedge2[i]-pedge[i+1]) so that the output is also intensive. By default, intensive=False is used and the input array is assumed to be extensive.
- edgelump (bool, default=True): If the max or min of pedge1 extend beyond the max or min of pedge2 and edgelump=True, the mass will be placed in in the first or last grid level of the output array2, ensuring that the mass of array2 is the same as array1. If edgelump=False, then the mass of array2 can be less than array1.
- Returns
- array2 (array (q,)): Array of remapped data. The length of array2 will be one less than pedge2.

def lon2i(lon, res=4): View Source

    315def lon2i(lon,res=4):
    316    '''Find I grid index for longitude
    317
    318    Parameters
    319    ----------
    320    lon : float
    321        longitude in degrees
    322    res : float, default=4
    323        resolution of grid. See `get_lon` for allowed values
    324
    325    Result
    326    ------
    327    index : int
    328        grid index location closest to longitude
    329    '''
    330    # Put longitude in expected GC the range [-180,180]
    331    alon = np.mod( np.asarray(lon) + 180, 360 ) - 180
    332    longrid = get_lon(res=res,center=True)
    333    nlon = len(longrid)
    334    f = interp1d(np.append(longrid,longrid[0]+360),
    335                np.append(np.arange(nlon),0),
    336                kind='nearest',
    337                fill_value='extrapolate')
    338    return f(alon).astype(np.int32)

Find I grid index for longitude

###### Parameters

- lon (float): longitude in degrees
- res (float, default=4): resolution of grid. See `get_lon` for allowed values

###### Result

index : int grid index location closest to longitude

def lat2j(lat, res=4): View Source

    340def lat2j(lat,res=4):
    341    '''Find J grid index for latitude
    342
    343    Parameters
    344    ----------
    345    lat : float
    346        latitude in degrees
    347    res : float, default=4
    348        resolution of grid. See `get_lon` for allowed values
    349
    350    Result
    351    ------
    352    index : int
    353        grid index location closest to latitude
    354    '''
    355    alat = np.asarray(lat)
    356    latgrid = get_lat(res=res,center=True)
    357    nlat = len(latgrid)
    358    if (np.any(alat<-90) or np.any(alat>90)):
    359        raise ValueError('lat must be in the range [-90,90]')
    360    f = interp1d(latgrid,
    361                np.arange(nlat),
    362                kind='nearest',
    363                fill_value='extrapolate')
    364    return f(alat).astype(np.int32)

Find J grid index for latitude

###### Parameters

- lat (float): latitude in degrees
- res (float, default=4): resolution of grid. See `get_lon` for allowed values

###### Result

index : int grid index location closest to latitude

def ll2ij(lon, lat, res=4): View Source

    366def ll2ij(lon,lat,res=4):
    367    '''Find (I,J) grid index for longitude, latitude
    368
    369
    370    Parameters
    371    ----------
    372    lon, lat : float
    373        longitude and latitude in degrees
    374    res : float, default=4
    375        resolution of grid. See `get_lon` for allowed values
    376
    377    Result
    378    ------
    379    i, j : int
    380        grid index location closest to longitude, latitude
    381    '''
    382    alon = np.asarray(lon)
    383    alat = np.asarray(lat)
    384    nlon = alon.size
    385    nlat = alat.size
    386    if (nlon==1 and nlat>1):
    387        alon = np.ones_like(alat) * alon
    388    elif (nlat==1 and nlon>1):
    389        alat = np.ones_like(alon) * alat
    390    elif nlon != nlat:
    391        raise ValueError('lon and lat must have the same length or be conformable')
    392    return lon2i(alon,res=res), lat2j(alat,res=res)

Find (I,J) grid index for longitude, latitude

###### Parameters

- lon, lat (float): longitude and latitude in degrees
- res (float, default=4): resolution of grid. See `get_lon` for allowed values

###### Result

i, j : int grid index location closest to longitude, latitude

def get_lon(res=4, center=None, edge=None): View Source

    394def get_lon(res=4, center=None, edge=None):
    395    '''Get array of longitudes for grid centers or edges.
    396
    397    GMAO grids assumed
    398    - res = 0.5 for 0.5x0.625
    399    - res = 2 for 2x2.5
    400    - res = 4 for 4x5
    401
    402    Parameters
    403    ----------
    404    res : float, default=4
    405        latitude resolution of grid
    406    center, edge : bool or None, default=None
    407        determines whether returned values are at grid center (center=True) or grid edge (edge=True)
    408        If both center and edge are None, then center is used
    409
    410    Returns
    411    -------
    412    lon : array (n,)
    413        longitude in degrees at grid center or edges
    414    '''
    415
    416    # Resolve any conflict between center and edge
    417    center, edge = _set_center_edge( center, edge, 'get_lon' )
    418
    419    # Longitude edges
    420    if res==2:
    421        dx=2.5
    422    elif res==4:
    423        dx=5
    424    elif res==0.5:
    425        dx=0.625
    426    else:
    427        raise NotImplementedError( 'Resolution '+str(res)+' is not defined' )
    428
    429    # Longitude edges
    430    lonedge = np.arange(-180 - dx/2, 180, dx)
    431
    432    # Get grid centers, if needed
    433    if center:
    434        lon = ( lonedge[0:-1] + lonedge[1:] ) / 2
    435    else:
    436        lon = lonedge
    437
    438    return lon

Get array of longitudes for grid centers or edges.

GMAO grids assumed

- res = 0.5 for 0.5x0.625
- res = 2 for 2x2.5
- res = 4 for 4x5

###### Parameters

- res (float, default=4): latitude resolution of grid
- center, edge (bool or None, default=None): determines whether returned values are at grid center (center=True) or grid edge (edge=True) If both center and edge are None, then center is used

###### Returns

- lon (array (n,)): longitude in degrees at grid center or edges

def get_lat(res=4, center=None, edge=None): View Source

    440def get_lat(res=4, center=None, edge=None):
    441    '''Get array of latitudes at grid center or edges
    442
    443    GMAO grids assumed
    444    - res = 0.5 for 0.5x0.625
    445    - res = 2 for 2x2.5
    446    - res = 4 for 4x5
    447
    448    Parameters
    449    ----------
    450    res : float, default=4
    451        latitude resolution of grid
    452    center, edge : bool or None, default=None
    453        determines whether returned values are at grid center (center=True) or grid edge (edge=True)
    454        If both center and edge are None, then center is used
    455
    456    Returns
    457    -------
    458    lat : array (n,)
    459        latitude in degrees at grid center or edges
    460    '''
    461
    462    # Resolve any conflict between center and edge
    463    center, edge = _set_center_edge( center, edge, 'get_lat' )
    464
    465    # Latitude edges
    466    if res==2:
    467        dy=2
    468    elif res==4:
    469        dy=4
    470    elif res==0.5:
    471        dy=0.5
    472    else:
    473        raise NotImplementedError( 'Resolution '+str(res)+' is not defined' )
    474
    475    # Latitude edges
    476    latedge = np.hstack([-90, np.arange(-90+dy/2,90,dy), 90] )
    477
    478    # Get grid centers, if needed
    479    if center:
    480        lat = ( latedge[0:-1] + latedge[1:] ) / 2
    481    else:
    482        lat = latedge
    483
    484    return lat

Get array of latitudes at grid center or edges

GMAO grids assumed

- res = 0.5 for 0.5x0.625
- res = 2 for 2x2.5
- res = 4 for 4x5

###### Parameters

- res (float, default=4): latitude resolution of grid
- center, edge (bool or None, default=None): determines whether returned values are at grid center (center=True) or grid edge (edge=True) If both center and edge are None, then center is used

###### Returns

- lat (array (n,)): latitude in degrees at grid center or edges

def get_lev_p(p_surf=1013.25, nlev=47, edge=False): View Source

    486def get_lev_p(p_surf=1013.25, nlev=47, edge=False ):
    487    '''Get array of pressures at grid centers or edges
    488
    489    Parameters
    490    ----------
    491    p_surf : float
    492        surface pressure, hPa
    493    nlev : {47,72}
    494        number of vertical levels (47 or 72)
    495    edge : bool
    496        specify edges (True) or centers (False, default)
    497
    498    Returns
    499    -------
    500    p : array of float
    501        pressures for each vertical level'''
    502
    503    # Get A,B values for vertical coordinate
    504    A, B = get_hybrid_ab(nlev=nlev, edge=edge )
    505
    506    # Pressure, hPa
    507    P = A + B * p_surf
    508
    509    return P

Get array of pressures at grid centers or edges

###### Parameters

- p_surf (float): surface pressure, hPa
- nlev ({47,72}): number of vertical levels (47 or 72)
- edge (bool): specify edges (True) or centers (False, default)

###### Returns

- p (array of float): pressures for each vertical level

def get_lev_z(TK=[273], p_surf=1013.25, p_edge=[0], Q=0, z_surf=0, edge=False): View Source

    512def get_lev_z(TK=[273], p_surf=1013.25, p_edge=[0], Q=0, z_surf=0, edge=False ):
    513    '''Calculate the geopotential heights of level edges or centers
    514
    515    Calculations use hypsometric equation, assuing hydrostatic balance
    516
    517    Parameters
    518    ----------
    519    TK : array of float
    520        Temperature for each grid level, K
    521    p_surf : float
    522        Surface pressure, hPa
    523    p_edge : array of float, optional
    524        Pressure at each grid level edge, hPa
    525        Will be calculated if length differs from TK
    526    Q : array of float
    527        Specific humidity for each grid level, kg/kg
    528    z_surf : float
    529        Surface altitude, m
    530    edge : bool
    531        specify edges (True) or centers (False, default) for output
    532
    533    Returns
    534    -------
    535    z : array of float
    536        Geopotential height for each grid edge or center, m
    537    '''
    538
    539    funname='geoschem: get_lev_z: '
    540
    541    # Acceleration due to gravity, m/s2
    542    g0 = 9.81
    543
    544    if len(TK)<=1:
    545        print(funname,'TK must be an array')
    546        return -1
    547
    548    # Number of levels in the temperature profile
    549    nlev = TK.size
    550
    551    # Calculate the pressure edges if they aren't provided in input
    552    if len(p_edge) != nlev+1:
    553        # Get pressure at level edges, hPa
    554        p_edge = get_lev_p(p_surf=p_surf, nlev=nlev, edge=True )
    555
    556    # Initialize array
    557    z_edge = np.zeros_like(p_edge)
    558    z_edge[0] = z_surf
    559
    560    # Virtual temperature, K
    561    Tv = TK * ( 1 + 0.61 * Q )
    562
    563    # Calculate altitudes with hypsometric equation, m
    564    for L in range(1, p_edge.size):
    565        z_edge[L] = z_edge[L-1] + 287 * Tv[L-1] / g0 * np.log( p_edge[L-1] / p_edge[L] )
    566
    567    # Choose grid edge or center altitudes, m
    568    if edge is True:
    569        z = z_edge
    570    else:
    571        z = (z_edge[0:-1] + z_edge[1:]) / 2
    572
    573    return z

Calculate the geopotential heights of level edges or centers

Calculations use hypsometric equation, assuing hydrostatic balance

###### Parameters

- TK (array of float): Temperature for each grid level, K
- p_surf (float): Surface pressure, hPa
- p_edge (array of float, optional): Pressure at each grid level edge, hPa Will be calculated if length differs from TK
- Q (array of float): Specific humidity for each grid level, kg/kg
- z_surf (float): Surface altitude, m
- edge (bool): specify edges (True) or centers (False, default) for output

###### Returns

- z (array of float): Geopotential height for each grid edge or center, m

def get_hybrid_ab(nlev=47, edge=False, center=False): View Source

    575def get_hybrid_ab(nlev=47, edge=False, center=False):
    576    '''Get A and B values for the GEOS-5/MERRA/MERRA2 hybrid sigma-pressure vertical coordinate.
    577
    578    The pressure at level edge L is P(L) = A(L) + B(L) * PS, where PS is surface pressure in hPa
    579
    580    Parameters
    581    ----------
    582    nlev : {47,72}
    583        number of levels in the desired grid (47 or 72)
    584    center, edge : bool
    585        specify if values for level center or edges are desired
    586
    587    Returns
    588    -------
    589    A, B : array (n,)
    590    '''
    591
    592    # Resolve any conflict between center and edge requests
    593    center, edge = _set_center_edge( center, edge, 'get_hybrid_ab')
    594
    595    if nlev == 47:
    596
    597        # A [hPa] for 47 levels (48 edges)
    598        A = np.array([ 0.000000E+00, 4.804826E-02, 6.593752E+00, 1.313480E+01,
    599               1.961311E+01, 2.609201E+01, 3.257081E+01, 3.898201E+01,
    600               4.533901E+01, 5.169611E+01, 5.805321E+01, 6.436264E+01,
    601               7.062198E+01, 7.883422E+01, 8.909992E+01, 9.936521E+01,
    602               1.091817E+02, 1.189586E+02, 1.286959E+02, 1.429100E+02,
    603               1.562600E+02, 1.696090E+02, 1.816190E+02, 1.930970E+02,
    604               2.032590E+02, 2.121500E+02, 2.187760E+02, 2.238980E+02,
    605               2.243630E+02, 2.168650E+02, 2.011920E+02, 1.769300E+02,
    606               1.503930E+02, 1.278370E+02, 1.086630E+02, 9.236572E+01,
    607               7.851231E+01, 5.638791E+01, 4.017541E+01, 2.836781E+01,
    608               1.979160E+01, 9.292942E+00, 4.076571E+00, 1.650790E+00,
    609               6.167791E-01, 2.113490E-01, 6.600001E-02, 1.000000E-02 ] )
    610
    611        # Bp [unitless] for 47 levels (48 edges)
    612        B = np.array([ 1.000000E+00, 9.849520E-01, 9.634060E-01, 9.418650E-01,
    613               9.203870E-01, 8.989080E-01, 8.774290E-01, 8.560180E-01,
    614               8.346609E-01, 8.133039E-01, 7.919469E-01, 7.706375E-01,
    615               7.493782E-01, 7.211660E-01, 6.858999E-01, 6.506349E-01,
    616               6.158184E-01, 5.810415E-01, 5.463042E-01, 4.945902E-01,
    617               4.437402E-01, 3.928911E-01, 3.433811E-01, 2.944031E-01,
    618               2.467411E-01, 2.003501E-01, 1.562241E-01, 1.136021E-01,
    619               6.372006E-02, 2.801004E-02, 6.960025E-03, 8.175413E-09,
    620               0.000000E+00, 0.000000E+00, 0.000000E+00, 0.000000E+00,
    621               0.000000E+00, 0.000000E+00, 0.000000E+00, 0.000000E+00,
    622               0.000000E+00, 0.000000E+00, 0.000000E+00, 0.000000E+00,
    623               0.000000E+00, 0.000000E+00, 0.000000E+00, 0.000000E+00 ] )
    624
    625    elif nlev == 72:
    626
    627        # A [hPa] for 72 levels (73 edges)
    628        A = np.array([ 0.000000E+00, 4.804826E-02, 6.593752E+00, 1.313480E+01,
    629               1.961311E+01, 2.609201E+01, 3.257081E+01, 3.898201E+01,
    630               4.533901E+01, 5.169611E+01, 5.805321E+01, 6.436264E+01,
    631               7.062198E+01, 7.883422E+01, 8.909992E+01, 9.936521E+01,
    632               1.091817E+02, 1.189586E+02, 1.286959E+02, 1.429100E+02,
    633               1.562600E+02, 1.696090E+02, 1.816190E+02, 1.930970E+02,
    634               2.032590E+02, 2.121500E+02, 2.187760E+02, 2.238980E+02,
    635               2.243630E+02, 2.168650E+02, 2.011920E+02, 1.769300E+02,
    636               1.503930E+02, 1.278370E+02, 1.086630E+02, 9.236572E+01,
    637               7.851231E+01, 6.660341E+01, 5.638791E+01, 4.764391E+01,
    638               4.017541E+01, 3.381001E+01, 2.836781E+01, 2.373041E+01,
    639               1.979160E+01, 1.645710E+01, 1.364340E+01, 1.127690E+01,
    640               9.292942E+00, 7.619842E+00, 6.216801E+00, 5.046801E+00,
    641               4.076571E+00, 3.276431E+00, 2.620211E+00, 2.084970E+00,
    642               1.650790E+00, 1.300510E+00, 1.019440E+00, 7.951341E-01,
    643               6.167791E-01, 4.758061E-01, 3.650411E-01, 2.785261E-01,
    644               2.113490E-01, 1.594950E-01, 1.197030E-01, 8.934502E-02,
    645               6.600001E-02, 4.758501E-02, 3.270000E-02, 2.000000E-02,
    646               1.000000E-02 ] )
    647
    648        # B [unitless] for 72 levels (73 edges)
    649        B = np.array([ 1.000000E+00, 9.849520E-01, 9.634060E-01, 9.418650E-01,
    650               9.203870E-01, 8.989080E-01, 8.774290E-01, 8.560180E-01,
    651               8.346609E-01, 8.133039E-01, 7.919469E-01, 7.706375E-01,
    652               7.493782E-01, 7.211660E-01, 6.858999E-01, 6.506349E-01,
    653               6.158184E-01, 5.810415E-01, 5.463042E-01, 4.945902E-01,
    654               4.437402E-01, 3.928911E-01, 3.433811E-01, 2.944031E-01,
    655               2.467411E-01, 2.003501E-01, 1.562241E-01, 1.136021E-01,
    656               6.372006E-02, 2.801004E-02, 6.960025E-03, 8.175413E-09,
    657               0.000000E+00, 0.000000E+00, 0.000000E+00, 0.000000E+00,
    658               0.000000E+00, 0.000000E+00, 0.000000E+00, 0.000000E+00,
    659               0.000000E+00, 0.000000E+00, 0.000000E+00, 0.000000E+00,
    660               0.000000E+00, 0.000000E+00, 0.000000E+00, 0.000000E+00,
    661               0.000000E+00, 0.000000E+00, 0.000000E+00, 0.000000E+00,
    662               0.000000E+00, 0.000000E+00, 0.000000E+00, 0.000000E+00,
    663               0.000000E+00, 0.000000E+00, 0.000000E+00, 0.000000E+00,
    664               0.000000E+00, 0.000000E+00, 0.000000E+00, 0.000000E+00,
    665               0.000000E+00, 0.000000E+00, 0.000000E+00, 0.000000E+00,
    666               0.000000E+00, 0.000000E+00, 0.000000E+00, 0.000000E+00,
    667               0.000000E+00 ])
    668
    669    else:
    670        raise ValueError( f'get_hybrid_ab: A, B not defined for nlev={nlev}' )
    671
    672    # Calculate values at level center
    673    if center is True:
    674
    675        A = (A[0:-1] + A[1:]) / 2
    676        B = (B[0:-1] + B[1:]) / 2
    677
    678    return A, B

Get A and B values for the GEOS-5/MERRA/MERRA2 hybrid sigma-pressure vertical coordinate.

The pressure at level edge L is P(L) = A(L) + B(L) * PS, where PS is surface pressure in hPa

###### Parameters

- nlev ({47,72}): number of levels in the desired grid (47 or 72)
- center, edge (bool): specify if values for level center or edges are desired

###### Returns

A, B : array (n,)