Python DataArray.attrs['long_name']方法代码示例

# 需要导入模块: from xarray import DataArray [as 别名]
# 或者: from xarray.DataArray import attrs['long_name'] [as 别名]

#.........这里部分代码省略.........
    For consistency with the analyses in the AMWG diagnostics and in papers
    such as Ming et el (2011), we should apply time averaging *before*
    calculating $ \Psi $.

    [zonal_mpsi_ncl]: https://www.ncl.ucar.edu/Document/Functions/Built-in/zonal_mpsi.shtml
    [zon_mpsi.f]: https://github.com/yyr/ncl/blob/34bafa4a78ba69ce8852212f59546bb433ce40c6/ni/src/lib/nfpfort/zon_mpsi.f
    [zmmsf.ncl]: http://www.ncl.ucar.edu/Applications/Scripts/zmmsf.ncl

    """

    #: Gravitational acceleration, m/s^2
    g = 9.80616

    #: Radius of earth, m
    a = 6.37122e6

    # Make sure the necessary fields are available in the dataset
    fields = ['V', 'PS']
    if diag_pressure: fields.extend(['hyam', 'hybm', 'P0'])
    for field in fields:
        if not (field in ds):
            raise KeyError('Expected to find key "%s" in dataset'
                           % field)

    # Transpose some dimensions (move to front) to facilitate
    # vertical integration in vector form
    v_wind = shuffle_dims(ds['V'], ['lev', 'lat'])
    ps = shuffle_dims(ds['PS'], ['lat', ]) # in Pa already

    levs = v_wind.lev.data
    lats = v_wind.lat.data
    nlev = len(levs)

    # 1) Construct the pressure thickness of each level, such that
    # pressure levels are mapped into the odd index and level interfaces
    # are mapped into the even indices
    if diag_pressure:
        pres = ( ds['hyam']*ds['P0'] + ds['hybm']*ds['PS'] )
        pres = shuffle_dims(pres, ['lev', 'lat'])
    else:
        pres = v_wind.lev*100. # hPa -> Pa
    # print(pres)

    # Mapping pressure levels and interfaces
    ptmp = np.zeros([2*nlev+1, ] + list(pres.shape[1:]))
    ptmp[0] = ptop # Pa
    ptmp[1::2] = pres.data
    ptmp[2:-1:2] = 0.5*(pres[1:].data + pres[:-1].data)
    ptmp[-1] = pbot # Pa

    # Compute thicknesses
    dp = np.zeros_like(ptmp)
    dp[1:-1:2] = (ptmp[2::2] - ptmp[0:-2:2]) # Pa

    # 2) Perform vertical integral of V as a function of level and
    # latitude, broadcasting over additional dimensions (time,
    # longitude, etc)
    psitmp = np.zeros([2*nlev+1, ] + list(v_wind.shape[1:]))
    vtmp = np.zeros([2*nlev+1, ] + list(v_wind.shape[1:]))

    psitmp[0] = 0.
    psitmp[1:2*nlev] = np.nan
    vtmp[0:-1:2] = 0.
    vtmp[1::2] = v_wind

    c = 2.*np.pi*a*np.cos(lats*np.pi/180.)/g # integral constant pre-factor

    # Integration loop, from top of atmosphere to bottom, using
    # a central difference method and 1/2-levels
    for klvl in range(1, 2*nlev, 2):

        integrand = c * np.transpose(vtmp[klvl]*dp[klvl])
        psitmp[klvl+1] = psitmp[klvl-1] - np.transpose(integrand)

    # Second loop to apply boundary condition to bottom of atmosphere, where
    # below-ground interfaces are changed to reflection points.
    bad_points = []
    for klvl in range(1, 2*nlev, 2):
        if np.any(ptmp[klvl] > ps):
            iinds, jinds = np.where(ptmp[klvl] > ps)
            # for k in ind_iter: print(k.shape)
            for (ilat, jlon) in zip(iinds, jinds):
                if (ilat, jlon) in bad_points: continue

                psitmp[klvl+1, ..., ilat, jlon] = \
                    -1.*psitmp[klvl-1, ..., ilat, jlon]
                psitmp[klvl+2:, ..., ilat, jlon] = np.nan
                bad_points.append((ilat, jlon))

    # psi on pressure levels = average of surrounding interfaces
    psitmp[1::2] = 0.5*(psitmp[2::2] + psitmp[0:-2:2])
    # Fix sign convention
    psitmp[1::2] *= -1.

    # Promote to a DataArray; attach metadata
    mpsi = DataArray(psitmp[1::2], coords=v_wind.coords, name='MPSI')
    mpsi.attrs['long_name'] = "Meridional Streamfunction"
    mpsi.attrs['units'] = 'kg/s'

    return mpsi