def test_weighted_sum_float32():
from esmlab.datasets import open_dataset
with open_dataset('cesm_pop_monthly') as ds:
weights = ds['TAREA'].astype(np.float32)
tmp_data = ds['TLAT']
tmp_data.values = np.where(np.greater_equal(ds['KMT'].values, 1),
ds['TLAT'], np.nan)
tmp_data = tmp_data.astype(np.float32)
w_sum = statistics.weighted_sum(tmp_data, dim=None, weights=weights)
assert tmp_data.attrs == w_sum.attrs
assert tmp_data.encoding == w_sum.encoding
def make_regional_timeseries(ds):
with open('pop_variable_defs.yml') as f:
pop_variable_defs = yaml.load(f)
# TODO: MASKED_AREA, MASKED_VOL could be precomputed and passed in...
rmask3d = region_mask_diagnostics(ds.isel(time=0))
MASKED_AREA, MASKED_VOL = make_masked_area_and_vol(ds.isel(time=0),
rmask3d,
ignore_ssh=True)
dso = xr.Dataset()
# copy time bounds
if 'bounds' in ds['time'].attrs:
tb_name = ds['time'].attrs['bounds']
elif 'time_bounds' in ds:
tb_name = 'time_bounds'
dso[tb_name] = ds[tb_name]
dz_150m = ds.dz[0:15].rename({'z_t': 'z_t_150m'})
variable_list = [v for v in ds.variables if 'time' in ds[v].dims]
for varname in variable_list:
if varname not in pop_variable_defs:
logging.warning(f'Missing definition of variable {varname}')
continue
info = pop_variable_defs[varname]
#-- get the variable
klevel = 0
#-- variable is a derived sum
if 'derived_sum' in info:
var = 0.
for v in info['derived_sum']:
var += ds[v]
var.name = varname
var.encoding = ds[info['derived_sum'][0]].encoding
#-- variable is derived by indexing "z_t"
elif 'derived_zindex' in info:
varname_from = info['derived_zindex'][0]
zlevel = info['derived_zindex'][1] * 100.
klevel = np.abs(ds.z_t.values - zlevel).argmin()
var = ds[varname_from].isel(z_t=klevel).drop(['z_t'])
var.name = varname
#-- variable is derived via "derived_var"
elif 'derived' in info:
var = derived_var(varname, ds)
#-- variable is read directly from file
else:
var = ds[varname]
#-- perform operations
if 'operations' in info:
for op in info['operations']:
if op == 'vertical_integral':
if var.dims[1] == 'z_t_150m':
var = gt.weighted_sum(var,
weights=dz_150m,
dim='z_t_150m')
elif var.dims[1] == 'z_t':
var = gt.weighted_sum(var, weights=ds.dz, dim='z_t')
else:
raise ValueError('Unknown operation: {0}'.format(op))
#-- compute averages
if info['category'] == 'tracer':
# add global inventory
dso[varname + '_avg'] = gt.weighted_mean(
var,
weights=MASKED_VOL.isel(region=0, drop=True),
dim=['z_t', 'nlat', 'nlon'],
apply_nan_mask=False)
dso[varname] = gt.weighted_mean(var,
weights=MASKED_AREA,
dim=['nlat', 'nlon'],
apply_nan_mask=False)
elif info['category'] in ['Cflux', 'Nflux', 'Siflux', 'Feflux']:
var = gt.weighted_sum(var,
weights=MASKED_AREA.isel(z_t=klevel).drop(
['z_t']),
dim=['nlat', 'nlon'])
dso[varname] = flux_weighted_area_sum_convert_units(
var, info['category'])
elif info['category'] in ['diagave']:
dso[varname] = gt.weighted_mean(
var,
weights=MASKED_AREA.isel(z_t=klevel).drop(['z_t']),
dim=['nlat', 'nlon'],
apply_nan_mask=False)
elif info['category'] in ['transport']:
dso[varname] = var
else:
raise ValueError('Unknown variable category {0}'.format(
info['category']))
return dso
def test_weighted_sum_ds(data, dim):
with pytest.warns(UserWarning):
w_sum = statistics.weighted_sum(ds, dim)
np.testing.assert_allclose(w_sum['variable_x'], ds['variable_x'].sum(dim))
def test_weighted_sum_da(data, dim, weights, axis):
with pytest.warns(UserWarning):
w_sum = statistics.weighted_sum(data, dim)
np.testing.assert_allclose(w_sum, data.sum(dim))
assert data.attrs == w_sum.attrs
assert data.encoding == w_sum.encoding
def test_weighted_sum():
with pytest.warns(UserWarning):
w_sum = statistics.weighted_sum(x)
np.testing.assert_allclose(w_sum, x.sum())
assert x.attrs == w_sum.attrs
assert x.encoding == w_sum.encoding