def test_grid_ops(all_datasets):
"""
Check that we get the same answer using Axis or Grid objects
"""
ds, periodic, expected = 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_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_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_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_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_wrong_axis_vector_input_axis(self):
ds, _, _ = datasets_grid_metric("C")
grid = Grid(ds)
msg = "Vector component with unknown axis provided. Grid has axes .*?"
with pytest.raises(
ValueError,
match=msg,
):
grid.diff({"wrong": xr.DataArray()}, "X")
def test_wrong_input_type_vector(self):
ds, _, _ = datasets_grid_metric("C")
grid = Grid(ds)
msg = "Dictionary inputs must have a DataArray as value. Got .*?"
with pytest.raises(
TypeError,
match=msg,
):
grid.diff({"X": "not_a_dataarray"}, "X")
def test_multiple_keys_vector_input(self):
ds, _, _ = datasets_grid_metric("C")
grid = Grid(ds)
msg = "Vector components provided as dictionaries should contain exactly one key/value pair. .*?"
with pytest.raises(
ValueError,
match=msg,
):
grid.diff({"X": xr.DataArray(), "Y": xr.DataArray()}, "X")
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_vector_missing_other_component(self, ds,
ds_face_connections_x_to_y):
grid = Grid(ds, face_connections=ds_face_connections_x_to_y)
msg = "Padding vector components requires `other_component` input"
with pytest.raises(ValueError, match=msg):
grid.diff(
{"X": ds.u},
"X",
other_component=None,
)
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_diff_interp_cubed_sphere(cs, cubed_sphere_connections):
grid = Grid(cs, face_connections=cubed_sphere_connections)
face, _ = xr.broadcast(cs.face, cs.data_c)
face_diff_x = grid.diff(face, "X")
np.testing.assert_allclose(face_diff_x[:, 0, 0], [-3, 1, 1, 1, 1, 2])
np.testing.assert_allclose(face_diff_x[:, -1, 0], [-3, 1, 1, 1, 1, 2])
face_diff_y = grid.diff(face, "Y")
np.testing.assert_allclose(face_diff_y[:, 0, 0], [-4, -3, -2, -1, 2, 5])
np.testing.assert_allclose(face_diff_y[:, 0, -1], [-4, -3, -2, -1, 2, 5])
def test_grid_no_coords(periodic_1d):
ds, periodic, expected = periodic_1d
ds_nocoords = ds.drop_dims(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_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_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_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):
"""Ensure that you can use xgcm with Xarray datasets that don't have dimension coordinates."""
ds, periodic, expected = periodic_1d
ds_nocoords = ds.drop_vars(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_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_grid_transform_noname(multidim_cases):
source, grid_kwargs, target, transform_kwargs, expected, error_flag = multidim_cases
axis = list(grid_kwargs["coords"].keys())[0]
grid = Grid(source, periodic=False, **grid_kwargs)
source_da = source.data
source_da.name = None
# the high level routines should be able to deal with all cases (no error flag exception like in the mid level)
transformed = grid.transform(source_da, axis, target, **transform_kwargs)
assert transformed.name is None
def test_grid_transform(all_cases):
source, grid_kwargs, target, transform_kwargs, expected, error_flag = all_cases
axis = list(grid_kwargs["coords"].keys())[0]
grid = Grid(source, periodic=False, **grid_kwargs)
# construct output name
transform_kwargs.setdefault("suffix", "")
output_name = "data" + transform_kwargs["suffix"]
# the high level routines should be able to deal with all cases (no error flag exception like in the mid level)
transformed = grid.transform(source.data, axis, target, **transform_kwargs)
xr.testing.assert_allclose(transformed, expected[output_name])
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_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_conservative_interp_warn():
(
source,
grid_kwargs,
target,
transform_kwargs,
_,
_,
) = construct_test_source_data(cases["conservative_depth_temp"])
axis = list(grid_kwargs["coords"].keys())[0]
grid = Grid(source, periodic=False, **grid_kwargs)
with pytest.warns(UserWarning):
_ = grid.transform(source.data, axis, target, **transform_kwargs)
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_get_metric_with_conditions_03a():
# 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"), "dx_n")
grid.set_metrics(("Y"), "dy_n")
get_metric = grid.get_metric(ds.v, ("X", "Y"))
metric_var_1 = ds.dx_n
metric_var_2 = ds.dy_n
expected_metric = (metric_var_1 * metric_var_2).reset_coords(drop=True)
xr.testing.assert_allclose(get_metric, expected_metric)
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_diff_interp_connected_grid_x_to_x(ds, ds_face_connections_x_to_x):
# simplest scenario with one face connection
grid = Grid(ds, face_connections=ds_face_connections_x_to_x)
diff_x = grid.diff(ds.data_c, "X", boundary="fill")
interp_x = grid.interp(ds.data_c, "X", boundary="fill")
# make sure the face connection got applied correctly
np.testing.assert_allclose(diff_x[1, :, 0],
ds.data_c[1, :, 0] - ds.data_c[0, :, -1])
np.testing.assert_allclose(
interp_x[1, :, 0], 0.5 * (ds.data_c[1, :, 0] + ds.data_c[0, :, -1]))
# make sure the left boundary got applied correctly
np.testing.assert_allclose(diff_x[0, :, 0], ds.data_c[0, :, 0] - 0.0)
np.testing.assert_allclose(interp_x[0, :, 0],
0.5 * (ds.data_c[0, :, 0] + 0.0))
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_chunking_dim_error():
"""Assure that error is raised when we chunk along the 'vertical' dimension"""
(
source,
grid_kwargs,
target,
transform_kwargs,
_,
_,
) = construct_test_source_data(cases["linear_depth_dens"])
source = source.chunk({"depth": 1})
axis = list(grid_kwargs["coords"].keys())[0]
grid = Grid(source, periodic=False, **grid_kwargs)
with pytest.raises(ValueError):
_ = grid.transform(source.data, axis, target, **transform_kwargs)
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)
