def xr_zonal_int_bins(da, LATS, AREA, dx=1):
""" integral over dx wide latitude bins
integrates da with AREA, then divides by width of zonal strip dx
input:
da .. 2D xr DataArray to be "zonally" integrated
LATS .. 2D xr DataArray latitude values of each cell
AREA .. 2D xr DataArray
dx .. width of latitude bands in degrees
lat_name .. xa/AREA coordinate name of the latitude variable
output:
xa_zonal_int .. 1D xr DataArray
lat centers can be accessed through xa_zonal_int.coords[f'{lat_name}_bins']
"""
assert type(da) == xr.core.dataarray.DataArray
assert type(AREA) == xr.core.dataarray.DataArray
assert np.shape(da)[-2:] == np.shape(AREA)
(z, lat, lon) = dll_from_arb_da(da)
lat_bins, lat_centers, lat_width = lat_binning(dx)
da_new = da * AREA
da_zonal_int = da_new.groupby_bins(LATS, lat_bins, labels=lat_centers).sum(
dim=f'stacked_{lat}_{lon}') / lat_width
return da_zonal_int
def xr_int_zonal(da, HTN, LATS, AREA, DZ):
""" integral along depth and zonal coordinates *[m^2] rectangular grid"""
shape = np.shape(da)
assert shape[-2:] == np.shape(HTN)[-2:]
assert shape[-2:] == np.shape(DZ)[-2:]
(z, lat, lon) = dll_from_arb_da(da)
if shape[-1] in [900, 320]: # rectangular `ocn_low` or `ocn_rect` grid
# print(np.shape(da))
# print(np.shape(HTN))
# print(np.shape(DZ))
int_zonal = (da * HTN * DZ).sum(dim=[z, lon]) # 1D (lat)
elif shape[-1] == 3600: # tripolar grid
int_vert = xr_int_vertical(da, DZ) # 2D
int_zonal = xr_zonal_int_bins(int_vert, LATS, AREA)
return int_zonal
def xr_int_zonal_level(da, HTN, LATS, AREA, DZ, dx=1):
""" zonal integrals for each level *[m] rectangular grid"""
(z, lat, lon) = dll_from_arb_da(da)
shape = np.shape(da)
assert shape[-2:] == np.shape(HTN)[-2:]
assert shape[-2:] == np.shape(DZ)[-2:]
if shape[-1] in [900, 320]: # rectangular grid
int_zonal_level = (da * HTN).sum(dim=[lon]) # 2D (z, lat)
elif shape[-1] == 3600: # tripolar grid
lat_bins, lat_centers, lat_width = lat_binning(dx)
km = len(da[z])
dz = DZ.max(dim=(lon, lat))
# construct new xr DataArray
assert 'time' in da.coords
lat_bin_name = f'TLAT_bins'
if da.coords['time'].size == 1: # single time files
array = np.zeros((km, len(lat_centers)))
coords = {z: da.coords[z], lat_bin_name: lat_centers}
int_zonal_level = xr.DataArray(data=array,
coords=coords,
dims=(z, lat_bin_name))
for k in range(km):
da_k = (da[k, :, :] * DZ[k, :, :]).drop('z_t')
int_zonal_level[k, :] = xr_zonal_int_bins(da_k, LATS,
AREA) / dz[k]
else:
array = np.zeros((da.coords['time'].size, km, len(lat_centers)))
coords = {
'time': da.coords['time'],
z: da.coords[z],
lat_bin_name: lat_centers
}
int_zonal_level = xr.DataArray(data=array,
coords=coords,
dims=('time', z, lat_bin_name))
for k in range(km):
da_k = (da[:, k, :, :] * DZ[k, :, :]).drop('z_t')
int_zonal_level[:, k, :] = xr_zonal_int_bins(da_k, LATS,
AREA) / dz[k]
return int_zonal_level
def xr_int_global(da, AREA, DZ):
""" global volume integral *[m^3] """
(z, lat, lon) = dll_from_arb_da(da)
return (da * AREA * DZ).sum(dim=[z, lat, lon]) # 0D
def xr_int_vertical(da, DZ):
""" vertical integral *[m] """
(z, lat, lon) = dll_from_arb_da(da)
return (da * DZ).sum(dim=z) # 2D (lat, lon)