def test_grid_dict_input_boundary_fill(nonperiodic_1d):
"""Test axis kwarg input functionality using dict input"""
ds, _, _ = nonperiodic_1d
grid_direct = Grid(ds, periodic=False, boundary="fill", fill_value=5)
grid_dict = Grid(ds, periodic=False, boundary={"X": "fill"}, fill_value={"X": 5})
assert grid_direct.axes["X"].fill_value == grid_dict.axes["X"].fill_value
assert grid_direct.axes["X"].boundary == grid_dict.axes["X"].boundary
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_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_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_invalid_fill_value_error():
ds = datasets["1d_left"]
with pytest.raises(ValueError):
Axis(ds, "X", fill_value="x")
with pytest.raises(ValueError):
Grid(ds, fill_value="bad")
with pytest.raises(ValueError):
Grid(ds, fill_value={"X": "bad"})
def test_boundary_kwarg_same_as_grid_constructor_kwarg():
ds = datasets["2d_left"]
grid1 = Grid(ds, periodic=False)
grid2 = Grid(ds, periodic=False, boundary={"X": "fill", "Y": "fill"})
actual1 = grid1.interp(ds.data_g, ("X", "Y"), boundary={"X": "fill", "Y": "fill"})
actual2 = grid2.interp(ds.data_g, ("X", "Y"))
xr.testing.assert_identical(actual1, actual2)
def test_dask_vs_eager(all_datasets, func, boundary):
ds, coords, metrics = datasets_grid_metric("C")
grid = Grid(ds, coords=coords)
grid_method = getattr(grid, func)
eager_result = grid_method(ds.tracer, "X", boundary=boundary)
ds = ds.chunk({"xt": 1, "yt": 1, "time": 1, "zt": 1})
grid = Grid(ds, coords=coords)
grid_method = getattr(grid, func)
dask_result = grid_method(ds.tracer, "X", boundary=boundary).compute()
xr.testing.assert_allclose(dask_result, eager_result)
def test_invalid_boundary_error():
ds = datasets["1d_left"]
with pytest.raises(ValueError):
Axis(ds, "X", boundary="bad")
with pytest.raises(ValueError):
Grid(ds, boundary="bad")
with pytest.raises(ValueError):
Grid(ds, boundary={"X": "bad"})
with pytest.raises(ValueError):
Grid(ds, boundary={"X": 0})
with pytest.raises(ValueError):
Grid(ds, boundary=0)
def test_grid_no_coords(periodic_1d):
ds, periodic, expected = periodic_1d
grid_expected = Grid(ds, periodic=periodic)
ds_nocoords = ds.drop(list(ds.dims.keys()))
coords = expected["axes"]
grid = Grid(ds_nocoords, periodic=periodic, coords=coords)
diff = grid.diff(ds["data_c"], "X")
assert len(diff.coords) == 0
interp = grid.interp(ds["data_c"], "X")
assert len(interp.coords) == 0
def test_create_grid_no_comodo(all_datasets):
ds, periodic, expected = all_datasets
grid_expected = Grid(ds, periodic=periodic)
ds_noattr = ds.copy()
for var in ds.variables:
ds_noattr[var].attrs.clear()
coords = expected["axes"]
grid = Grid(ds_noattr, periodic=periodic, coords=coords)
for axis_name_expected in grid_expected.axes:
axis_expected = grid_expected.axes[axis_name_expected]
axis_actual = grid.axes[axis_name_expected]
_assert_axes_equal(axis_expected, axis_actual)
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_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_vector_connected_grid_x_to_y(ds, ds_face_connections_x_to_y,
boundary):
# one face connection, rotated
grid = Grid(
ds,
face_connections=ds_face_connections_x_to_y,
boundary=boundary,
fill_value=1,
periodic=False,
)
# TODO: Remove the periodic once that is deprecated.
# ! Set boundary on grid, so it is applied to all axes.
# TODO: modify the non velocity tests too (after release)
# modify the values of the dataset, so we know what to expect from the output
# TODO: Maybe change this in the dataset definition?
u_modifier = xr.DataArray([-2, -1], dims="face")
v_modifier = xr.DataArray([1, 1], dims="face")
u = ds.u * 0 + u_modifier
v = ds.v * 0 + v_modifier
# no need to check for diff vs interp. They all go through the same dispatch
# v is the interesting variable here because it involves a sign change for this
# connection (see https://github.com/xgcm/xgcm/issues/410#issue-1098348557)
v_out = grid.interp({"Y": v}, "X", other_component={"X": u})
# the test case is set up in a way that all interpolated values for u should be 1
# if the face connection is done properly
np.testing.assert_allclose(v_out.data, 1)
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_keep_coords(funcname, gridtype):
ds, coords, metrics = datasets_grid_metric(gridtype)
ds = ds.assign_coords(yt_bis=ds["yt"], xt_bis=ds["xt"])
grid = Grid(ds, coords=coords, metrics=metrics)
func = getattr(grid, funcname)
for axis_name in grid.axes.keys():
result = func(ds.tracer, axis_name)
base_coords = list(result.dims)
augmented_coords = [
c for c in ds.coords
if set(ds[c].dims).issubset(result.dims) and c not in result.dims
]
if funcname in ["integrate", "average"]:
assert set(result.coords) == set(base_coords + augmented_coords)
else:
assert set(result.coords) == set(base_coords)
# TODO: why is the behavior different for integrate and average?
if funcname not in ["integrate", "average"]:
result = func(ds.tracer, axis_name, keep_coords=False)
assert set(result.coords) == set(base_coords)
result = func(ds.tracer, axis_name, keep_coords=True)
assert set(result.coords) == set(base_coords + augmented_coords)
def test_grid_ops(all_datasets):
"""
Check that we get the same answer using Axis or Grid objects
"""
ds, periodic, _ = all_datasets
grid = Grid(ds, periodic=periodic)
for axis_name in grid.axes.keys():
try:
ax_periodic = axis_name in periodic
except TypeError:
ax_periodic = periodic
axis = Axis(ds, axis_name, periodic=ax_periodic)
bcs = [None] if ax_periodic else ["fill", "extend"]
for varname in ["data_c", "data_g"]:
for boundary in bcs:
da_interp = grid.interp(ds[varname],
axis_name,
boundary=boundary)
da_interp_ax = axis.interp(ds[varname], boundary=boundary)
assert da_interp.equals(da_interp_ax)
da_diff = grid.diff(ds[varname], axis_name, boundary=boundary)
da_diff_ax = axis.diff(ds[varname], boundary=boundary)
assert da_diff.equals(da_diff_ax)
if boundary is not None:
da_cumsum = grid.cumsum(ds[varname],
axis_name,
boundary=boundary)
da_cumsum_ax = axis.cumsum(ds[varname], boundary=boundary)
assert da_cumsum.equals(da_cumsum_ax)
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_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_grid_transform_bypass_checks(bypass_checks):
"""Check that the bypass checks option still delivers the right results for monotonically increasing data"""
(
source,
grid_kwargs,
target,
transform_kwargs,
expected,
_,
) = construct_test_source_data(cases["linear_depth_dens"])
axis = list(grid_kwargs["coords"].keys())[0]
grid = Grid(source, periodic=False, **grid_kwargs)
target_data = transform_kwargs.pop("target_data", None)
transformed = grid.transform(
source.data,
axis,
target,
target_data=target_data,
bypass_checks=bypass_checks,
**transform_kwargs
)
xr.testing.assert_allclose(transformed, expected.data)
def test_grid_transform_noname_targetdata():
"""Check handling of a `target_data` input without name"""
(
source,
grid_kwargs,
target,
transform_kwargs,
_,
_,
) = construct_test_source_data(cases["linear_depth_dens"])
axis = list(grid_kwargs["coords"].keys())[0]
grid = Grid(source, periodic=False, **grid_kwargs)
source_da = source.data
target_data = transform_kwargs.pop("target_data")
target_data.name = None
# the name of target_data is only used if `target` is provided as numpy array
target = target.data
# the high level routines should be able to deal with all cases (no error flag exception like in the mid level)
with pytest.warns(UserWarning):
transformed = grid.transform(source_da,
axis,
target,
target_data=target_data,
**transform_kwargs)
"TRANSFORMED_DIMENSION" in transformed.dims
def test_grid_transform_multidim_other_dims_error(request, multidim_cases):
# broadcast the 1d column agains some other dims and make sure that the 1d results are still valid
source, grid_kwargs, target, transform_kwargs, expected, error_flag = multidim_cases
na = 3
source = source * xr.DataArray(np.ones([na]), dims=["a"])
# broadcast the target, but in this case
# rename one of the dimensions of the target array, which is not along the
# axis of transformation (this could be the case if e.g. temperature is on a different
# x grid than velocity)
target_data = transform_kwargs.pop("target_data", None)
if target_data is not None:
target_data = target_data * xr.DataArray(np.ones([na]),
dims=["a_other"])
# calculate the multidimensional result
axis = list(grid_kwargs["coords"].keys())[0]
grid = Grid(source, periodic=False, **grid_kwargs)
with pytest.raises(ValueError):
_ = grid.transform(source.data,
axis,
target,
target_data=target_data,
**transform_kwargs)
else:
# When target_data is none its taken as a 1D-coordinate, no checking needed
pytest.skip()
def test_grid_transform_input_check():
(
source,
grid_kwargs,
target,
transform_kwargs,
_,
_,
) = construct_test_source_data(cases["linear_depth_dens"])
axis = list(grid_kwargs["coords"].keys())[0]
grid = Grid(source, periodic=False, **grid_kwargs)
# construct output name
transform_kwargs.setdefault("suffix", "")
# Make sure that a sensible error is raised if xr.Dataset is provided
# for either one of `source`, `target` or `target_data` input arguments.
match_msg = r"needs to be a"
with pytest.raises(ValueError, match=r"`da` " + match_msg):
grid.transform(source, axis, target, **transform_kwargs)
with pytest.raises(ValueError, match=match_msg):
grid.transform(source.data, axis, target.to_dataset(name="dummy"),
**transform_kwargs)
transform_kwargs["target_data"] = transform_kwargs[
"target_data"].to_dataset(name="dummy")
with pytest.raises(ValueError, match=match_msg):
grid.transform(source.data, axis, target, **transform_kwargs)
def test_grid_transform_multidim(request, client, multidim_cases):
# broadcast the 1d column agains some other dims and make sure that the 1d results are still valid
source, grid_kwargs, target, transform_kwargs, expected, error_flag = multidim_cases
na = 8
source = source.expand_dims(a=na).copy(deep=True)
# broadcast the target_data manually
target_data = transform_kwargs.pop("target_data", None)
if target_data is not None:
target_data = target_data.expand_dims(a=na).copy(deep=True)
if client != "no_client":
target_data = target_data.chunk({"a": 1})
if client != "no_client":
source = source.chunk({"a": 1})
# calculate the multidimensional result
axis = list(grid_kwargs["coords"].keys())[0]
grid = Grid(source, periodic=False, **grid_kwargs)
# the high level tests should deal with all error cases
client = request.getfixturevalue(client)
transformed = grid.transform(source.data,
axis,
target,
target_data=target_data,
**transform_kwargs).load()
_, expected_broadcasted = xr.broadcast(transformed, expected)
xr.testing.assert_allclose(transformed, expected_broadcasted.data)
def test_keep_coords_deprecation():
ds, coords, metrics = datasets_grid_metric("B")
ds = ds.assign_coords(yt_bis=ds["yt"], xt_bis=ds["xt"])
grid = Grid(ds, coords=coords, metrics=metrics)
for axis_name in grid.axes.keys():
with pytest.warns(DeprecationWarning):
grid.diff(ds.tracer, axis_name, keep_coords=False)
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_wrong_input_type_scalar(self):
ds, _, _ = datasets_grid_metric("C")
grid = Grid(ds)
msg = "All data arguments must be either a DataArray or Dictionary .*?"
with pytest.raises(
TypeError,
match=msg,
):
grid.apply_as_grid_ufunc(lambda x: x, "not_a_dataarray", "X")
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_transform_error_periodic(multidim_cases):
source, grid_kwargs, target, transform_kwargs, expected, error_flag = multidim_cases
axis = list(grid_kwargs["coords"].keys())[0]
grid = Grid(source, **grid_kwargs)
with pytest.raises(ValueError):
_ = grid.transform(source.data, axis, target, **transform_kwargs)
def test_input_dim_notfound():
data = np.random.rand(4, 5)
coord = np.random.rand(4, 5)
ds = xr.DataArray(data, dims=["x", "y"], coords={
"c": (["x", "y"], coord)
}).to_dataset(name="data")
msg = r"Could not find dimension `other` \(for the `center` position on axis `X`\) in input dataset."
with pytest.raises(ValueError, match=msg):
Grid(ds, coords={"X": {"center": "other"}})
def test_input_not_dims():
data = np.random.rand(4, 5)
coord = np.random.rand(4, 5)
ds = xr.DataArray(data, dims=["x", "y"], coords={
"c": (["x", "y"], coord)
}).to_dataset(name="data")
msg = r"is not a dimension in the input dataset"
with pytest.raises(ValueError, match=msg):
Grid(ds, coords={"X": {"center": "c"}})
