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_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_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_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_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_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_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_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_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_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_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)
