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_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_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_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_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_interp_like(
metric_axes, metric_name, periodic, boundary, boundary_expected, fill_value
):
ds, coords, _ = datasets_grid_metric("C")
grid = Grid(ds, coords=coords, periodic=periodic)
grid.set_metrics(metric_axes, metric_name)
metric_available = grid._metrics.get(frozenset(metric_axes), None)
metric_available = metric_available[0]
interp_metric = grid.interp_like(
metric_available, ds.u, boundary=boundary, fill_value=fill_value
)
expected_metric = grid.interp(
ds[metric_name], metric_axes, boundary=boundary_expected, fill_value=fill_value
)
xr.testing.assert_allclose(interp_metric, expected_metric)
