def test_DerivedMapping_dataset():
derived_state = DerivedMapping(ds)
keys = ["T", "q", "cos_zenith_angle"]
ds_derived_state = derived_state.dataset(keys)
assert isinstance(ds_derived_state, xr.Dataset)
for existing_var in ["T", "q"]:
np.testing.assert_array_almost_equal(
ds_derived_state[existing_var], ds[existing_var]
)
def add_derived_data(variables: Sequence[str], ds: xr.Dataset) -> xr.Dataset:
"""
Overlay the DerivedMapping and grab a dataset of specified variables
Args:
variables: All variables (derived and non-derived) to include in the
dataset.
"""
derived_mapping = DerivedMapping(ds)
return derived_mapping.dataset(variables)