def test_1d_padded_but_no_change_in_grid_position(self):
def diff_center_to_center_second_order(a):
b = a[..., 2:]
c = a[..., :-2]
return 0.5 * (b - c)
grid = create_1d_test_grid("depth")
da = grid._ds.depth_c**2
da.coords["depth_c"] = grid._ds.depth_c
diffed = 0.5 * (da - da.roll(depth_c=2, roll_coords=False)).data
expected = xr.DataArray(
diffed, dims=["depth_c"], coords={"depth_c": grid._ds.depth_c}
)
def pad_args(func, pad_width):
def padding_version_of_func(*args):
padded_args = [
np.pad(a, pad_width=pad_width, mode="wrap") for a in args
]
return func(*padded_args)
return padding_version_of_func
# Test direct application
result = apply_as_grid_ufunc(
pad_args(diff_center_to_center_second_order, pad_width=[(2, 0)]),
da,
axis=[("depth",)],
grid=grid,
signature="(X:center)->(X:center)",
boundary_width=None,
)
assert_equal(result, expected)
def test_zero_width_boundary(self):
def increment(x):
"""Mocking up a function which can only act on in-memory arrays, and requires no padding"""
if isinstance(x, np.ndarray):
return np.add(x, 1)
else:
raise TypeError
grid = create_1d_test_grid("depth")
a = np.sin(grid._ds.depth_g * 2 * np.pi / 9).chunk(2)
a.coords["depth_g"] = grid._ds.depth_g
expected = a + 1
result = apply_as_grid_ufunc(
increment,
a,
axis=[("depth",)],
grid=grid,
signature="(X:left)->(X:left)",
boundary_width=None,
dask="allowed",
map_overlap=True,
).compute()
assert_equal(result, expected)
# in this case the result should be the same as using just map_blocks
expected_data = dask.array.map_blocks(increment, a.data)
np.testing.assert_equal(result.data, expected_data)
def test_cumsum(self):
def cumsum_center_to_left(a):
return np.cumsum(a, axis=-1)[..., :-1]
grid = create_1d_test_grid("depth")
da = grid._ds.depth_c**2
da.coords["depth_c"] = grid._ds.depth_c
cum = da.cumsum(dim="depth_c").roll(depth_c=1, roll_coords=False)
cum[0] = 0
expected = xr.DataArray(
cum.data, dims=["depth_g"], coords={"depth_g": grid._ds.depth_g}
)
result = apply_as_grid_ufunc(
cumsum_center_to_left,
da,
axis=[("depth",)],
grid=grid,
signature="(X:center)->(X:left)",
boundary_width={"X": (1, 0)},
boundary="fill",
fill_value=0,
pad_before_func=False,
)
assert_equal(result, expected)
def test_apply_along_one_axis(self):
grid = create_2d_test_grid("lon", "lat")
def diff_center_to_left(a):
return a - np.roll(a, shift=-1, axis=-1)
da = grid._ds.lat_c**2 + grid._ds.lon_c**2
diffed = (da - da.roll(lon_c=-1, roll_coords=False)).data
expected = xr.DataArray(
diffed,
dims=["lat_c", "lon_g"],
coords={"lat_c": grid._ds.lat_c, "lon_g": grid._ds.lon_g},
)
# Test direct application
result = apply_as_grid_ufunc(
diff_center_to_left,
da,
axis=[("lon",)],
grid=grid,
signature="(X:center)->(X:left)",
)
assert_equal(result, expected)
# Test decorator
@as_grid_ufunc()
def diff_center_to_left(
a: Annotated[np.ndarray, "X:center"]
) -> Annotated[np.ndarray, "X:left"]:
return a - np.roll(a, shift=-1, axis=-1)
result = diff_center_to_left(grid, da, axis=[("lon",)])
assert_equal(result, expected)
def test_1d_unchanging_size_but_padded_dask_parallelized(self):
"""
This test checks that the process of padding a non-chunked core dimension doesn't turn it into a chunked core
dimension. See GH #430.
"""
def diff_center_to_left(a):
return a[..., 1:] - a[..., :-1]
grid = create_1d_test_grid("depth")
da = np.sin(grid._ds.depth_c * 2 * np.pi / 9).chunk()
da.coords["depth_c"] = grid._ds.depth_c
diffed = (da - da.roll(depth_c=1, roll_coords=False)).data
expected = xr.DataArray(
diffed, dims=["depth_g"], coords={"depth_g": grid._ds.depth_g}
).compute()
# Test direct application
result = apply_as_grid_ufunc(
diff_center_to_left,
da,
axis=[("depth",)],
grid=grid,
signature="(X:center)->(X:left)",
boundary_width={"X": (1, 0)},
dask="parallelized",
).compute()
assert_equal(result, expected)
# Test Grid method
result = grid.apply_as_grid_ufunc(
diff_center_to_left,
da,
axis=[("depth",)],
signature="(X:center)->(X:left)",
boundary_width={"X": (1, 0)},
dask="parallelized",
)
assert_equal(result, expected)
# Test decorator
@as_grid_ufunc(
"(X:center)->(X:left)",
boundary_width={"X": (1, 0)},
dask="parallelized",
)
def diff_center_to_left(a):
return a[..., 1:] - a[..., :-1]
result = diff_center_to_left(
grid,
da,
axis=[("depth",)],
).compute()
assert_equal(result, expected)
def test_chunked_core_dims_unchanging_chunksize(self):
def diff_center_to_left(a):
return a[..., 1:] - a[..., :-1]
grid = create_1d_test_grid("depth")
da = np.sin(grid._ds.depth_c * 2 * np.pi / 9).chunk(3)
da.coords["depth_c"] = grid._ds.depth_c
diffed = (da - da.roll(depth_c=1, roll_coords=False)).data
expected = xr.DataArray(
diffed, dims=["depth_g"], coords={"depth_g": grid._ds.depth_g}
).compute()
# Test direct application
result = apply_as_grid_ufunc(
diff_center_to_left,
da,
axis=[("depth",)],
grid=grid,
signature="(X:center)->(X:left)",
boundary_width={"X": (1, 0)},
dask="allowed",
map_overlap=True,
).compute()
assert_equal(result, expected)
# Test Grid method
result = grid.apply_as_grid_ufunc(
diff_center_to_left,
da,
axis=[("depth",)],
signature="(X:center)->(X:left)",
boundary_width={"X": (1, 0)},
dask="allowed",
map_overlap=True,
)
assert_equal(result, expected)
# Test decorator
@as_grid_ufunc(
boundary_width={"X": (1, 0)},
dask="allowed",
map_overlap=True,
)
def diff_center_to_left(
a: Annotated[np.ndarray, "X:center"]
) -> Annotated[np.ndarray, "X:left"]:
return a[..., 1:] - a[..., :-1]
result = diff_center_to_left(
grid,
da,
axis=[("depth",)],
).compute()
assert_equal(result, expected)
def test_1d_overlap_dask_allowed(self):
from dask.array import map_overlap
def diff_center_to_left(a):
return a - np.roll(a, shift=-1)
def diff_overlap(a):
return map_overlap(diff_center_to_left, a, depth=1, boundary="periodic")
grid = create_1d_test_grid("depth")
da = np.sin(grid._ds.depth_c * 2 * np.pi / 9).chunk(1)
da.coords["depth_c"] = grid._ds.depth_c
diffed = (da - da.roll(depth_c=-1, roll_coords=False)).data
expected = xr.DataArray(
diffed, dims=["depth_g"], coords={"depth_g": grid._ds.depth_g}
).compute()
# Test direct application
result = apply_as_grid_ufunc(
diff_overlap,
da,
axis=[("depth",)],
grid=grid,
signature="(X:center)->(X:left)",
dask="allowed",
).compute()
assert_equal(result, expected)
# Test Grid method
result = grid.apply_as_grid_ufunc(
diff_overlap,
da,
axis=[("depth",)],
signature="(X:center)->(X:left)",
dask="allowed",
)
assert_equal(result, expected)
# Test decorator
@as_grid_ufunc(dask="allowed")
def diff_overlap(
a: Annotated[np.ndarray, "X:center"]
) -> Annotated[np.ndarray, "X:left"]:
return map_overlap(diff_center_to_left, a, depth=1, boundary="periodic")
result = diff_overlap(
grid,
da,
axis=[("depth",)],
).compute()
assert_equal(result, expected)
def test_multiple_outputs(self):
def diff_center_to_inner(a, axis):
result = a - np.roll(a, shift=1, axis=axis)
return np.delete(result, 0, axis) # remove first element along axis
def grad_to_inner(a):
return diff_center_to_inner(a, axis=0), diff_center_to_inner(a, axis=1)
grid = create_2d_test_grid("lon", "lat")
a = grid._ds.lon_c**2 + grid._ds.lat_c**2
expected_u = 2 * grid._ds.lon_i.expand_dims(dim={"lat_c": len(grid._ds.lat_c)})
expected_u.coords["lat_c"] = grid._ds.lat_c
expected_v = 2 * grid._ds.lat_i.expand_dims(dim={"lon_c": len(grid._ds.lon_c)})
expected_v.coords["lon_c"] = grid._ds.lon_c
# Test direct application
u, v = apply_as_grid_ufunc(
grad_to_inner,
a,
axis=[("lon", "lat")],
grid=grid,
signature="(X:center,Y:center)->(X:inner,Y:center),(X:center,Y:inner)",
)
assert_equal(u.T, expected_u)
assert_equal(v, expected_v)
# Test Grid method
u, v = grid.apply_as_grid_ufunc(
grad_to_inner,
a,
axis=[("lon", "lat")],
signature="(X:center,Y:center)->(X:inner,Y:center),(X:center,Y:inner)",
)
assert_equal(u.T, expected_u)
assert_equal(v, expected_v)
# Test decorator
@as_grid_ufunc()
def grad_to_inner(
a: Annotated[np.ndarray, "X:center,Y:center"]
) -> Tuple[
Annotated[np.ndarray, "X:inner,Y:center"],
Annotated[np.ndarray, "X:center,Y:inner"],
]:
return diff_center_to_inner(a, axis=0), diff_center_to_inner(a, axis=1)
u, v = grad_to_inner(grid, a, axis=[("lon", "lat")])
assert_equal(u.T, expected_u)
assert_equal(v, expected_v)
def test_multiple_outputs(self):
def diff_center_to_inner(a, axis):
result = a - np.roll(a, shift=1, axis=axis)
return np.delete(result, 0, axis) # remove first element along axis
def grad_to_inner(a):
return diff_center_to_inner(a, axis=0), diff_center_to_inner(a, axis=1)
grid = create_2d_test_grid("lon", "lat")
a = (grid._ds.lon_c**2 + grid._ds.lat_c**2).chunk(1)
# Test direct application
with pytest.raises(NotImplementedError, match="multiple outputs"):
apply_as_grid_ufunc(
grad_to_inner,
a,
axis=[("lon", "lat")],
grid=grid,
signature="(X:center,Y:center)->(X:inner,Y:center),(X:center,Y:inner)",
map_overlap=True,
dask="allowed",
)
def test_input_on_wrong_positions(self):
grid = create_1d_test_grid("depth")
grid._ds.drop_vars("depth_o")
da = np.sin(grid._ds.depth_g * 2 * np.pi / 9)
with pytest.raises(
ValueError,
match=re.escape("Axis:positions pair depth:outer does not exist"),
):
da: Annotated[np.ndarray, "X:outer"]
apply_as_grid_ufunc(
lambda x: x, da, axis=[("depth",)], grid=grid, signature="(X:outer)->()"
)
with pytest.raises(ValueError, match="coordinate depth_c does not appear"):
da: Annotated[np.ndarray, "X:center"]
apply_as_grid_ufunc(
lambda x: x,
da,
axis=[("depth",)],
grid=grid,
signature="(X:center)->()",
)
def test_1d_changing_size_dask_parallelized(self):
def interp_center_to_inner(a):
return 0.5 * (a[:-1] + a[1:])
grid = create_1d_test_grid("depth")
da = xr.DataArray(
np.arange(10, 19), dims=["depth_c"], coords={"depth_c": grid._ds.depth_c}
).chunk()
expected = da.interp(depth_c=np.arange(1.5, 9), method="linear").rename(
depth_c="depth_i"
)
# Test direct application
result = apply_as_grid_ufunc(
interp_center_to_inner,
da,
axis=[("depth",)],
grid=grid,
signature="(X:center)->(X:inner)",
dask="parallelized",
).compute()
assert_equal(result, expected)
# Test Grid method
result = grid.apply_as_grid_ufunc(
interp_center_to_inner,
da,
axis=[("depth",)],
signature="(X:center)->(X:inner)",
dask="parallelized",
).compute()
assert_equal(result, expected)
# Test decorator
@as_grid_ufunc(dask="parallelized")
def interp_center_to_inner(
a: Annotated[np.ndarray, "X:center"]
) -> Annotated[np.ndarray, "X:inner"]:
return 0.5 * (a[:-1] + a[1:])
result = interp_center_to_inner(grid, da, axis=[("depth",)]).compute()
assert_equal(result, expected)
def test_multiple_inputs(self):
def inner_product_left_right(a, b):
return np.inner(a, b)
grid = create_1d_test_grid("depth")
a = np.sin(grid._ds.depth_g * 2 * np.pi / 9)
a.coords["depth_g"] = grid._ds.depth_g
b = np.cos(grid._ds.depth_r * 2 * np.pi / 9)
b.coords["depth_r"] = grid._ds.depth_r
expected = xr.DataArray(np.inner(a, b))
# Test direct application
result = apply_as_grid_ufunc(
inner_product_left_right,
a,
b,
axis=[("depth",), ("depth",)],
grid=grid,
signature="(X:left),(X:right)->()",
)
assert_equal(result, expected)
# Test Grid method
result = grid.apply_as_grid_ufunc(
inner_product_left_right,
a,
b,
axis=[("depth",), ("depth",)],
signature="(X:left),(X:right)->()",
)
assert_equal(result, expected)
# Test decorator
@as_grid_ufunc()
def inner_product_left_right(
a: Annotated[np.ndarray, "X:left"], b: Annotated[np.ndarray, "X:right"]
):
return np.inner(a, b)
result = inner_product_left_right(grid, a, b, axis=[("depth",), ("depth",)])
assert_equal(result, expected)
def test_1d_unchanging_size_no_dask(self):
def diff_center_to_left(a):
return a - np.roll(a, shift=-1)
grid = create_1d_test_grid("depth")
da = np.sin(grid._ds.depth_c * 2 * np.pi / 9)
da.coords["depth_c"] = grid._ds.depth_c
diffed = (da - da.roll(depth_c=-1, roll_coords=False)).data
expected = xr.DataArray(
diffed, dims=["depth_g"], coords={"depth_g": grid._ds.depth_g}
)
# Test direct application
result = apply_as_grid_ufunc(
diff_center_to_left,
da,
axis=[("depth",)],
grid=grid,
signature="(X:center)->(X:left)",
)
assert_equal(result, expected)
# Test Grid method
result = grid.apply_as_grid_ufunc(
diff_center_to_left, da, axis=[("depth",)], signature="(X:center)->(X:left)"
)
assert_equal(result, expected)
# Test decorator
@as_grid_ufunc()
def diff_center_to_left(
a: Annotated[np.ndarray, "X:center"]
) -> Annotated[np.ndarray, "X:left"]:
return a - np.roll(a, shift=-1)
result = diff_center_to_left(grid, da, axis=[("depth",)])
assert_equal(result, expected)
def test_1d_padded_but_no_change_in_grid_position(self):
def diff_center_to_center_second_order(a):
return 0.5 * (a[..., 2:] - a[..., :-2])
grid = create_1d_test_grid("depth")
da = np.sin(grid._ds.depth_c * 2 * np.pi / 9)
da.coords["depth_c"] = grid._ds.depth_c
diffed = 0.5 * (da - da.roll(depth_c=2, roll_coords=False)).data
expected = xr.DataArray(
diffed, dims=["depth_c"], coords={"depth_c": grid._ds.depth_c}
)
# Test direct application
result = apply_as_grid_ufunc(
diff_center_to_center_second_order,
da,
axis=[("depth",)],
grid=grid,
signature="(X:center)->(X:center)",
boundary_width={"X": (2, 0)},
)
assert_equal(result, expected)