def test_multiple_metrics_per_axis():
# copied from test_derivatives.py - should refactor
dx = 10.0
ds = xr.Dataset(
{
"foo": (("XC", ), [1.0, 2.0, 4.0, 3.0]),
"bar": (("XG", ), [10.0, 20.0, 30.0, 40.0]),
},
coords={
"XC": (("XC", ), [0.5, 1.5, 2.5, 3.5]),
"XG": (("XG", ), [0, 1.0, 2.0, 3.0]),
"dXC": (("XC", ), [dx, dx, dx, dx]),
"dXG": (("XG", ), [dx, dx, dx, dx]),
},
)
grid = Grid(
ds,
coords={"X": {
"center": "XC",
"left": "XG"
}},
metrics={("X", ): ["dXC", "dXG"]},
periodic=True,
)
assert grid.get_metric(ds.foo,
("X", )).equals(ds.dXC.reset_coords(drop=True))
assert grid.get_metric(ds.bar,
("X", )).equals(ds.dXG.reset_coords(drop=True))
def test_get_metric(axes, data_var, drop_vars, metric_expected_list,
expected_error):
ds_full = datasets()
ds = ds_full["C"]
metrics = ds_full["metrics"]
# drop metrics according to drop_vars input, and remove from metrics input
if drop_vars:
print(drop_vars)
ds = ds.drop(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=ds_full["coords"], metrics=metrics)
if expected_error:
with pytest.raises(expected_error):
metric = grid.get_metric(ds[data_var], axes)
else:
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_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_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_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_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_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_metrics_2d_grid():
# this is a uniform grid
# a non-uniform grid would provide a more rigorous test
dx = 10.0
dy = 11.0
area = 120.0
ny, nx = 7, 9
ds = xr.Dataset(
{"foo": (("YC", "XC"), np.ones((ny, nx)))},
coords={
"XC": (("XC",), np.arange(nx)),
"dX": (("XC",), np.full(nx, dx)),
"YC": (("YC",), np.arange(ny)),
"dY": (("YC",), np.full(ny, dy)),
"area": (("YC", "XC"), np.full((ny, nx), area)),
},
)
grid = Grid(
ds,
coords={"X": {"center": "XC"}, "Y": {"center": "YC"}},
metrics={("X",): ["dX"], ("Y",): ["dY"], ("X", "Y"): ["area"]},
)
assert ds.dX.reset_coords(drop=True).equals(grid.get_metric(ds.foo, ("X",)))
assert ds.dY.reset_coords(drop=True).equals(grid.get_metric(ds.foo, ("Y",)))
assert ds.area.reset_coords(drop=True).equals(grid.get_metric(ds.foo, ("X", "Y")))
assert ds.area.reset_coords(drop=True).equals(grid.get_metric(ds.foo, ("Y", "X")))
# try with no area metric:
grid = Grid(
ds,
coords={"X": {"center": "XC"}, "Y": {"center": "YC"}},
metrics={("X",): ["dX"], ("Y",): ["dY"]},
)
dxdy = (ds.dX * ds.dY).reset_coords(drop=True).transpose("YC", "XC")
actual = grid.get_metric(ds.foo, ("Y", "X")).transpose("YC", "XC")
assert dxdy.equals(actual)
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)
