def test_derivative_b_grid(self):
# test derivatives with synthetic B grid data
ds, coords, metrics = datasets_grid_metric("B")
grid = Grid(ds, coords=coords, metrics=metrics)
# tracer point
var = "tracer"
test_axes = ["X", "Y", "Z"]
test_dx = ["dx_e", "dy_n", "dz_w"]
for ax, dx in zip(test_axes, test_dx):
_run_single_derivative_test(grid, ax, ds[var], ds[dx])
# zonal velocity point
var = "u"
test_dx = ["dx_n", "dy_e", "dz_w_ne"]
for ax, dx in zip(test_axes, test_dx):
_run_single_derivative_test(grid, ax, ds[var], ds[dx])
# meridional velocity point
var = "v"
test_dx = ["dx_n", "dy_e", "dz_w_ne"]
for ax, dx in zip(test_axes, test_dx):
_run_single_derivative_test(grid, ax, ds[var], ds[dx])
# vertical velocity point
var = "wt"
test_dx = ["dx_e", "dy_n", "dz_t"]
for ax, dx in zip(test_axes, test_dx):
_run_single_derivative_test(grid, ax, ds[var], ds[dx])
def test_set_metric():
ds, coords, metrics = datasets_grid_metric("C")
expected_metrics = {k: [ds[va] for va in v] for k, v in metrics.items()}
grid = Grid(ds, coords=coords, metrics=metrics)
grid_manual = Grid(ds, coords=coords)
for key, value in metrics.items():
grid_manual.set_metrics(key, value)
assert len(grid._metrics) > 0
for k, v in expected_metrics.items():
k = frozenset(k)
assert k in grid._metrics.keys()
assert k in grid_manual._metrics.keys()
for metric_expected, metric in zip(v, grid_manual._metrics[k]):
xr.testing.assert_allclose(metric_expected.reset_coords(drop=True), metric)
for metric_expected, metric in zip(v, grid._metrics[k]):
xr.testing.assert_allclose(metric_expected.reset_coords(drop=True), metric)
def test_padding_fill(self, boundary_width, fill_value):
ds, coords, _ = datasets_grid_metric("C")
grid = Grid(ds, coords=coords)
data = ds.tracer
# iterate over all axes
expected = data.copy(deep=True)
for ax, widths in boundary_width.items():
dim = grid._get_dims_from_axis(
data, ax)[0] # ? not entirely sure why this is a list
expected = expected.pad({dim: widths},
"constant",
constant_values=fill_value)
# we intentionally strip all coords from padded arrays
expected = _strip_all_coords(expected)
result = pad(
data,
grid,
boundary="fill",
boundary_width=boundary_width,
fill_value=fill_value,
other_component=None,
)
xr.testing.assert_allclose(expected, result)
def test_set_metric_overwrite_true(
metric_axes, exist_metric_varname, add_metric_varname, expected_varname
):
ds, coords, metrics = datasets_grid_metric("C")
ds = ds.assign_coords({add_metric_varname[0]: ds[exist_metric_varname[1]] * 10})
metrics = {
k: [m for m in v if m in exist_metric_varname] for k, v in metrics.items()
}
grid = Grid(ds, coords=coords, metrics=metrics)
for av in add_metric_varname:
grid.set_metrics(metric_axes, av, overwrite=True)
key = frozenset(list(metric_axes))
set_metric = grid._metrics.get(key)
expected_metric = []
for ev in expected_varname:
metric_var = ds[ev].reset_coords(drop=True)
expected_metric.append(metric_var)
assert len(set_metric) == len(expected_metric)
for i in range(len(set_metric)):
assert set_metric[i].equals(expected_metric[i])
def test_boundary_global_input(funcname, boundary, fill_value):
"""Test that globally defined boundary values result in
the same output as when the parameters are defined on either
the grid or axis methods
"""
ds, coords, metrics = datasets_grid_metric("C")
axis = "X"
# Test results by globally specifying fill value/boundary on grid object
grid_global = Grid(
ds,
coords=coords,
metrics=metrics,
periodic=False,
boundary=boundary,
fill_value=fill_value,
)
func_global = getattr(grid_global, funcname)
global_result = func_global(ds.tracer, axis)
# Test results by manually specifying fill value/boundary on grid method
grid_manual = Grid(
ds, coords=coords, metrics=metrics, periodic=False, boundary=boundary
)
func_manual = getattr(grid_manual, funcname)
manual_result = func_manual(
ds.tracer, axis, boundary=boundary, fill_value=fill_value
)
xr.testing.assert_allclose(global_result, manual_result)
def test_padding_mixed(self, boundary_width):
ds, coords, _ = datasets_grid_metric("C")
grid = Grid(ds, coords=coords)
data = ds.tracer
axis_padding_mapping = {"X": "periodic", "Y": "extend"}
method_mapping = {
"periodic": "wrap",
"extend": "edge",
}
# iterate over all axes
expected = data.copy(deep=True)
for ax, widths in boundary_width.items():
dim = grid._get_dims_from_axis(
data, ax)[0] # ? not entirely sure why this is a list
expected = expected.pad({dim: widths},
method_mapping[axis_padding_mapping[ax]])
# we intentionally strip all coords from padded arrays
expected = _strip_all_coords(expected)
result = pad(
data,
grid,
boundary=axis_padding_mapping,
boundary_width=boundary_width,
fill_value=None,
other_component=None,
)
xr.testing.assert_allclose(expected, result)
def test_weighted_metric_multi_axis(self, funcname, grid_type, variable,
multi_axis, metric_weighted, boundary):
"""tests if the output for multiple axis is the same as when
executing the single axis ops in serial"""
ds, coords, metrics = datasets_grid_metric(grid_type)
grid = Grid(ds, coords=coords, metrics=metrics)
func = getattr(grid, funcname)
expected = ds[variable]
for ax in multi_axis:
if isinstance(metric_weighted, dict):
metric_weighted_axis = metric_weighted[ax]
else:
metric_weighted_axis = metric_weighted
expected = func(
expected,
ax,
metric_weighted=metric_weighted_axis,
boundary=boundary,
)
new = func(
ds[variable],
multi_axis,
metric_weighted=metric_weighted,
boundary=boundary,
)
assert new.equals(expected)
def test_missingaxis(axis, funcname):
# Error should be raised if application axes
# include dimension not in datasets
ds, coords, metrics = datasets_grid_metric("C")
del coords[axis]
del_metrics = [k for k in metrics.keys() if axis in k]
for dm in del_metrics:
del metrics[dm]
grid = Grid(ds, coords=coords, metrics=metrics)
func = getattr(grid, funcname)
if funcname == "cumint":
# cumint needs a boundary...
kwargs = dict(boundary="fill")
else:
kwargs = dict()
match_message = (
"Unable to find any combinations of metrics for array dims.*%s.*" %
axis)
with pytest.raises(KeyError, match=match_message):
func(ds.tracer, ["X", "Y", "Z"])
with pytest.raises(KeyError, match=match_message):
func(ds, axis, **kwargs)
if axis == "Y":
# test two missing axes at the same time
del coords["X"]
del_metrics = [k for k in metrics.keys() if "X" in k]
for dm in del_metrics:
del metrics[dm]
grid = Grid(ds, coords=coords, metrics=metrics)
func = getattr(grid, funcname)
if funcname == "cumint":
# cumint needs a boundary...
kwargs = dict(boundary="fill")
else:
kwargs = dict()
match_message = (
"Unable to find any combinations of metrics for array dims.*X.*Y.*Z.*"
)
with pytest.raises(KeyError, match=match_message):
func(ds, ["X", "Y", "Z"], **kwargs)
match_message = (
"Unable to find any combinations of metrics for array dims.*X.*Y.*"
)
with pytest.raises(KeyError, match=match_message):
func(ds, ("X", "Y"), **kwargs)
def test_get_metric_with_conditions_01():
# Condition 1: metric with matching axes and dimensions exist
ds, coords, metrics = datasets_grid_metric("C")
grid = Grid(ds, coords=coords, metrics=metrics)
get_metric = grid.get_metric(ds.v, ("X", "Y"))
expected_metric = ds["area_n"].reset_coords(drop=True)
xr.testing.assert_allclose(get_metric, expected_metric)
def test_set_metric_key_errors(metric_axes, add_metric):
ds, coords, metrics = datasets_grid_metric("C")
grid = Grid(ds, coords=coords, metrics=metrics)
if len(metric_axes) == 1:
with pytest.raises(KeyError, match="not found in dataset."):
grid.set_metrics(metric_axes, add_metric)
else:
with pytest.raises(KeyError, match="not compatible with grid axes"):
grid.set_metrics(metric_axes, add_metric)
def test_set_metric_value_errors(metric_axes, overwrite_metric, add_metric):
ds, coords, metrics = datasets_grid_metric("C")
if add_metric is not None:
ds = ds.assign_coords({overwrite_metric: ds[add_metric] * 10})
grid = Grid(ds, coords=coords, metrics=metrics)
with pytest.raises(ValueError, match="setting overwrite=True."):
grid.set_metrics(metric_axes, overwrite_metric)
def test_interp_conservative(grid_type, variable, axis, metric_weighted):
ds, coords, metrics = datasets_grid_metric(grid_type)
grid = Grid(ds, coords=coords, metrics=metrics)
metric = grid.get_metric(ds[variable], metric_weighted)
expected_raw = grid.interp(ds[variable] * metric, axis)
metric_new = grid.get_metric(expected_raw, metric_weighted)
expected = expected_raw / metric_new
new = grid.interp(ds[variable], axis, metric_weighted=metric_weighted)
assert new.equals(expected)
def test_get_metric_with_conditions_02a(periodic):
# Condition 2, case a: interpolate metric with matching axis to desired dimensions
ds, coords, _ = datasets_grid_metric("C")
grid = Grid(ds, coords=coords, periodic=periodic, boundary="extend")
grid.set_metrics(("X", "Y"), "area_e")
get_metric = grid.get_metric(ds.v, ("X", "Y"))
expected_metric = grid.interp(ds["area_e"], ("X", "Y"))
xr.testing.assert_allclose(get_metric, expected_metric)
def test_get_metric_orig(axes, data_var, drop_vars, metric_expected_list):
ds, coords, metrics = datasets_grid_metric("C")
# drop metrics according to drop_vars input, and remove from metrics input
if drop_vars:
ds = ds.drop_vars(drop_vars)
metrics = {k: [a for a in v if a not in drop_vars] for k, v in metrics.items()}
grid = Grid(ds, coords=coords, metrics=metrics)
metric = grid.get_metric(ds[data_var], axes)
expected_metrics = [ds[me].reset_coords(drop=True) for me in metric_expected_list]
expected = functools.reduce(operator.mul, expected_metrics, 1)
assert metric.equals(expected)
def test_cgrid(self, funcname, boundary, periodic):
ds, coords, metrics = datasets_grid_metric("C")
grid = Grid(ds, coords=coords, metrics=metrics, periodic=periodic)
if funcname == "cumint":
# cumint needs a boundary
kwargs = dict(boundary=boundary)
else:
# integrate and average don't use the boundary input
kwargs = dict()
func = getattr(grid, funcname)
# test tracer position
for axis, metric_name, dim in zip(
["X", "Y", "Z", ["X", "Y"], ["X", "Y", "Z"]],
["dx_t", "dy_t", "dz_t", "area_t", "volume_t"],
["xt", "yt", "zt", ["xt", "yt"], ["xt", "yt", "zt"]],
):
new = func(ds.tracer, axis, **kwargs)
expected = _expected_result(
ds.tracer, ds[metric_name], grid, dim, axis, funcname, **kwargs
)
assert_allclose(new, expected)
# test with tuple input if list is provided
if isinstance(axis, list):
new = func(ds.tracer, tuple(axis), **kwargs)
assert_allclose(new, expected)
# test u positon
for axis, metric_name, dim in zip(
["X", "Y", ["X", "Y"]],
["dx_e", "dy_e", "area_e"], # need more metrics?
["xu", "yt", ["xu", "yt"]],
):
new = func(ds.u, axis, **kwargs)
expected = _expected_result(
ds.u, ds[metric_name], grid, dim, axis, funcname, **kwargs
)
assert_allclose(new, expected)
# test v positon
for axis, metric_name, dim in zip(
["X", "Y", ["X", "Y"]],
["dx_n", "dy_n", "area_n"], # need more metrics?
["xt", "yu", ["xt", "yu"]],
):
new = func(ds.v, axis, **kwargs)
expected = _expected_result(
ds.v, ds[metric_name], grid, dim, axis, funcname, **kwargs
)
assert_allclose(new, expected)
def test_weighted_metric(
self, funcname, grid_type, variable, axis, metric_weighted, periodic, boundary
):
"""tests the correct execution of weighted ops along a single axis"""
# metric_weighted allows the interpolation of e.g. a surface flux to be conservative
# It multiplies the values with a metric like the area, then performs interpolation
# and divides by the same metric (area) for the new grid position
ds, coords, metrics = datasets_grid_metric(grid_type)
grid = Grid(ds, coords=coords, metrics=metrics, periodic=periodic)
func = getattr(grid, funcname)
metric = grid.get_metric(ds[variable], metric_weighted)
expected_raw = func(ds[variable] * metric, axis, boundary=boundary)
metric_new = grid.get_metric(expected_raw, metric_weighted)
expected = expected_raw / metric_new
new = func(
ds[variable], axis, metric_weighted=metric_weighted, boundary=boundary
)
assert new.equals(expected)
@pytest.mark.parametrize(
"multi_axis", ["X", ["X"], ("Y"), ["X", "Y"], ("Y", "X")]
)
def test_weighted_metric_multi_axis(
self, funcname, grid_type, variable, multi_axis, metric_weighted, boundary
):
"""tests if the output for multiple axis is the same as when
executing the single axis ops in serial"""
ds, coords, metrics = datasets_grid_metric(grid_type)
grid = Grid(ds, coords=coords, metrics=metrics)
func = getattr(grid, funcname)
expected = ds[variable]
for ax in multi_axis:
if isinstance(metric_weighted, dict):
metric_weighted_axis = metric_weighted[ax]
else:
metric_weighted_axis = metric_weighted
expected = func(
expected,
ax,
metric_weighted=metric_weighted_axis,
boundary=boundary,
)
new = func(
ds[variable],
multi_axis,
metric_weighted=metric_weighted,
boundary=boundary,
)
assert new.equals(expected)
def test_get_metric_with_conditions_04a():
# Condition 4, case a: 1 metric on the wrong position (must interpolate before multiplying)
ds, coords, _ = datasets_grid_metric("C")
grid = Grid(ds, coords=coords)
grid.set_metrics(("X"), "dx_t")
grid.set_metrics(("Y"), "dy_n")
get_metric = grid.get_metric(ds.v, ("X", "Y"))
interp_metric = grid.interp(ds.dx_t, "Y")
expected_metric = (interp_metric * ds.dy_n).reset_coords(drop=True)
xr.testing.assert_allclose(get_metric, expected_metric)
def test_get_metric_with_conditions_03b():
# Condition 3: use provided metrics with matching dimensions to calculate for required metric
ds, coords, metrics = datasets_grid_metric("C")
grid = Grid(ds, coords=coords)
grid.set_metrics(("X", "Y"), "area_t")
grid.set_metrics(("Z"), "dz_t")
get_metric = grid.get_metric(ds.tracer, ("X", "Y", "Z"))
metric_var_1 = ds.area_t
metric_var_2 = ds.dz_t
expected_metric = (metric_var_1 * metric_var_2).reset_coords(drop=True)
xr.testing.assert_allclose(get_metric, expected_metric)
def test_get_metric_with_conditions_02b():
# Condition 2, case b: get_metric should select for the metric with matching axes and interpolate from there,
# even if other metrics in the desired positions are available
ds, coords, _ = datasets_grid_metric("C")
grid = Grid(ds, coords=coords)
grid.set_metrics(("X", "Y"), "area_e")
grid.set_metrics(("X"), "dx_n")
grid.set_metrics(("Y"), "dx_n")
get_metric = grid.get_metric(ds.v, ("X", "Y"))
expected_metric = grid.interp(ds["area_e"], ("X", "Y"))
xr.testing.assert_allclose(get_metric, expected_metric)
def test_bgrid(funcname):
ds, coords, metrics = datasets_grid_metric("B")
grid = Grid(ds, coords=coords, metrics=metrics)
if funcname == "cumint":
# cumint needs a boundary...
kwargs = dict(boundary="fill")
else:
kwargs = dict()
func = getattr(grid, funcname)
# test tracer position
for axis, metric_name, dim in zip(
["X", "Y", "Z", ["X", "Y"], ["X", "Y", "Z"]],
["dx_t", "dy_t", "dz_t", "area_t", "volume_t"],
["xt", "yt", "zt", ["xt", "yt"], ["xt", "yt", "zt"]],
):
new = func(ds.tracer, axis, **kwargs)
expected = _expected_result(ds.tracer, ds[metric_name], grid, dim,
axis, funcname)
assert new.equals(expected)
# test with tuple input if list is provided
if isinstance(axis, list):
new = func(ds.tracer, tuple(axis), **kwargs)
assert new.equals(expected)
# test u position
for axis, metric_name, dim in zip(
["X", "Y", ["X", "Y"]],
["dx_ne", "dy_ne", "area_ne"], # need more metrics?
["xu", "yu", ["xu", "yu"]],
):
new = func(ds.u, axis, **kwargs)
expected = _expected_result(ds.u, ds[metric_name], grid, dim, axis,
funcname)
assert new.equals(expected)
# test v position
for axis, metric_name, dim in zip(
["X", "Y", ["X", "Y"]],
["dx_ne", "dy_ne", "area_ne"], # need more metrics?
["xu", "yu", ["xu", "yu"]],
):
new = func(ds.v, axis, **kwargs)
expected = _expected_result(ds.v, ds[metric_name], grid, dim, axis,
funcname)
assert new.equals(expected)
def test_missingaxis(self, axis, funcname, periodic, boundary):
# Error should be raised if application axes include dimension not in datasets
ds, coords, metrics = datasets_grid_metric("C")
del coords[axis]
del_metrics = [k for k in metrics.keys() if axis in k]
for dm in del_metrics:
del metrics[dm]
grid = Grid(ds, coords=coords, metrics=metrics, periodic=periodic)
func = getattr(grid, funcname)
if funcname == "cumint":
# cumint needs a boundary
kwargs = dict(boundary=boundary)
else:
kwargs = dict()
with pytest.raises(KeyError, match="Did not find axis"):
func(ds.tracer, ["X", "Y", "Z"], **kwargs)
if axis == "Y":
# test two missing axes at the same time
del coords["X"]
del_metrics = [k for k in metrics.keys() if "X" in k]
for dm in del_metrics:
del metrics[dm]
grid = Grid(ds, coords=coords, metrics=metrics)
func = getattr(grid, funcname)
if funcname == "cumint":
# cumint needs a boundary
kwargs = dict(boundary="fill")
else:
kwargs = dict()
with pytest.raises(KeyError, match="Did not find axis"):
func(ds.tracer, ["X", "Y", "Z"], **kwargs)
with pytest.raises(KeyError, match="Did not find axis"):
func(ds.tracer, ("X", "Y"), **kwargs)
def test_metric_axes_missing_from_array(self, funcname, periodic, boundary):
ds, coords, metrics = datasets_grid_metric("C")
grid = Grid(ds, coords=coords, metrics=metrics, periodic=periodic)
if funcname == "cumint":
# cumint needs a boundary
kwargs = dict(boundary="fill")
else:
kwargs = dict()
func = getattr(grid, funcname)
with pytest.raises(ValueError, match="Did not find single matching dimension"):
func(ds.tracer.mean("xt"), "X", **kwargs)
with pytest.raises(ValueError, match="Did not find single matching dimension"):
func(ds.tracer.mean("xt"), ["X", "Y", "Z"], **kwargs)
def test_weighted_metric(self, funcname, grid_type, variable, axis,
metric_weighted, periodic, boundary):
"""tests the correct execution of weighted ops along a single axis"""
# metric_weighted allows the interpolation of e.g. a surface flux to be conservative
# It multiplies the values with a metric like the area, then performs interpolation
# and divides by the same metric (area) for the new grid position
ds, coords, metrics = datasets_grid_metric(grid_type)
grid = Grid(ds, coords=coords, metrics=metrics, periodic=periodic)
func = getattr(grid, funcname)
metric = grid.get_metric(ds[variable], metric_weighted)
expected_raw = func(ds[variable] * metric, axis, boundary=boundary)
metric_new = grid.get_metric(expected_raw, metric_weighted)
expected = expected_raw / metric_new
new = func(ds[variable],
axis,
metric_weighted=metric_weighted,
boundary=boundary)
assert new.equals(expected)
def test_missingdim(funcname):
ds, coords, metrics = datasets_grid_metric("C")
grid = Grid(ds, coords=coords, metrics=metrics)
func = getattr(grid, funcname)
if funcname == "cumint":
# cumint needs a boundary...
kwargs = dict(boundary="fill")
else:
kwargs = dict()
match_message = "Unable to find any combinations of metrics for array dims.*X.*"
with pytest.raises(KeyError, match=match_message):
func(ds.tracer.mean("xt"), "X", **kwargs)
match_message = (
"Unable to find any combinations of metrics for array dims.*X.*Y.*Z.*")
with pytest.raises(KeyError, match=match_message):
func(ds.tracer.mean("xt"), ["X", "Y", "Z"])
def test_assign_metric(key, metric_vars):
ds, coords, _ = datasets_grid_metric("C")
_ = Grid(ds, coords=coords, metrics={key: metric_vars})
