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_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_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_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.diff("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_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_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_derivative_uniform_grid(self):
# this is a uniform grid
# a non-uniform grid would provide a more rigorous test
dx = 10.0
dy = 10.0
arr = [
[1.0, 2.0, 4.0, 3.0],
[4.0, 7.0, 1.0, 2.0],
[3.0, 1.0, 0.0, 9.0],
[8.0, 5.0, 2.0, 1.0],
]
ds = xr.Dataset(
{"foo": (("XC", "YC"), arr)},
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]),
"YC": (("YC", ), [0.5, 1.5, 2.5, 3.5]),
"YG": (("YG", ), [0, 1.0, 2.0, 3.0]),
"dYC": (("YC", ), [dy, dy, dy, dy]),
"dYG": (("YG", ), [dy, dy, dy, dy]),
},
)
grid = Grid(
ds,
coords={
"X": {
"center": "XC",
"left": "XG"
},
"Y": {
"center": "YC",
"left": "YG"
},
},
metrics={
("X", ): ["dXC", "dXG"],
("Y", ): ["dYC", "dYG"]
},
periodic=True,
)
# Test x direction
dfoo_dx = grid.derivative(ds.foo, "X")
expected = grid.diff(ds.foo, "X") / dx
assert dfoo_dx.equals(expected)
# Test y direction
dfoo_dy = grid.derivative(ds.foo, "Y")
expected = grid.diff(ds.foo, "Y") / dy
assert dfoo_dy.equals(expected)
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_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_diff_interp_connected_grid_x_to_y(ds, ds_face_connections_x_to_y):
# one face connection, rotated
grid = Grid(ds, face_connections=ds_face_connections_x_to_y)
diff_x = grid.diff(ds.data_c, 'X', boundary='fill')
interp_x = grid.interp(ds.data_c, 'X', boundary='fill')
diff_y = grid.diff(ds.data_c, 'Y', boundary='fill')
interp_y = grid.interp(ds.data_c, 'Y', boundary='fill')
# make sure the face connection got applied correctly
# non-same axis connections require rotation
# ravel everything to avoid dealing with broadcasting
np.testing.assert_allclose(
diff_y.data[1, 0, :].ravel(),
ds.data_c.data[1, 0, :].ravel() - ds.data_c.data[0, ::-1, -1].ravel())
np.testing.assert_allclose(
interp_y.data[1, 0, :].ravel(),
0.5 * (ds.data_c.data[1, 0, :].ravel() +
ds.data_c.data[0, ::-1, -1].ravel()))
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_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_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_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_derivative_uniform_grid():
# this is a uniform grid
# a non-uniform grid would provide a more rigorous test
dx = 10.0
ds = xr.Dataset(
{"foo": (("XC",), [1.0, 2.0, 4.0, 3.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,
)
dfoo_dx = grid.derivative(ds.foo, "X")
expected = grid.diff(ds.foo, "X") / dx
assert dfoo_dx.equals(expected)
