Skip to content

Commit

Permalink
Allow non scalar dx/dy coords in structured regridding; fix potential…
Browse files Browse the repository at this point in the history
… y-flipping in regridding structured -> unstructured.

Fixes #275
Fixes #157
  • Loading branch information
Huite committed Aug 13, 2024
1 parent 42f10ab commit d242307
Show file tree
Hide file tree
Showing 3 changed files with 77 additions and 6 deletions.
38 changes: 38 additions & 0 deletions tests/test_regrid/test_regridder.py
Original file line number Diff line number Diff line change
Expand Up @@ -283,3 +283,41 @@ def test_create_percentile_method():
weights = np.ones_like(values)
workspace = np.zeros_like(values)
assert median(values, weights, workspace) == 2


def test_directional_dependence():
# Increasing / decreasing x or y direction shouldn't matter for the result.
da = xr.DataArray(
data=[[1.0, 2.0], [3.0, 4.0]],
coords={"y": [17.5, 12.5], "x": [2.5, 7.5]},
dims=("y", "x"),
)
target_da = xr.DataArray(
data=[[np.nan, np.nan], [np.nan, np.nan]],
coords={"y": [10.0, 20.0], "x": [0.0, 10.0]},
dims=("y", "x"),
)
target_uda = xu.UgridDataArray.from_structured(target_da)

flip = slice(None, None, -1)
flipy = da.isel(y=flip)
flipx = da.isel(x=flip)
flipxy = da.isel(x=flip, y=flip)
uda = xu.UgridDataArray.from_structured(da)
uda_flipxy = xu.UgridDataArray.from_structured(flipxy)

# Structured target: test whether the result is the same regardless of source
# orientation.
result = []
for source in [da, flipy, flipx, flipxy, uda, uda_flipxy]:
regridder = xu.OverlapRegridder(source, taget=target_da)
result.append(regridder.regrid(source))
assert all(result[0] == item for item in result[1:])

# Unstructured target: test whether the result is the same regardless of
# source orientation.
result = []
for source in [da, flipy, flipx, flipxy, uda, uda_flipxy]:
regridder = xu.OverlapRegridder(source, taget=target_uda)
result.append(regridder.regrid(source))
assert all(result[0] == item for item in result[1:])
23 changes: 23 additions & 0 deletions tests/test_regrid/test_structured.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import numpy as np
import pytest
import xarray as xr

from xugrid.regrid.structured import StructuredGrid1d, StructuredGrid2d

Expand Down Expand Up @@ -404,3 +405,25 @@ def test_linear_weights_2d(
np.array([5, 5, 5, 5, 6, 6, 6, 6, 9, 9, 9, 9, 10, 10, 10, 10]),
np.array([1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0]),
)


def test_nonscalar_dx():
da = xr.DataArray(
[1, 2, 3], coords={"x": [1, 2, 3], "dx": ("x", [1, 1, 1])}, dims=("x",)
)
grid = StructuredGrid1d(da, name="x")
actual = xr.DataArray([1, 2, 3], coords=grid.coords, dims=grid.dims)
assert actual.identical(da)


def test_directional_bounds():
da = xr.DataArray([1, 2, 3], coords={"y": [1, 2, 3]}, dims=("y",))
decreasing = da.isel(y=slice(None, None, -1))
grid_inc = StructuredGrid1d(da, name="y")
grid_dec = StructuredGrid1d(decreasing, name="y")
assert grid_inc.flipped is False
assert grid_dec.flipped is True
assert np.array_equal(grid_inc.bounds, grid_dec.bounds)
assert np.array_equal(
grid_inc.directional_bounds, grid_dec.directional_bounds[::-1]
)
22 changes: 16 additions & 6 deletions xugrid/regrid/structured.py
Original file line number Diff line number Diff line change
Expand Up @@ -83,10 +83,12 @@ def __init__(self, obj: Union[xr.DataArray, xr.Dataset], name: str):

@property
def coords(self) -> dict:
return {
self.name: self.index,
self.dname: self.dvalue,
}
coords = {self.name: self.index}
if isinstance(self.dvalue, np.ndarray):
coords[self.dname] = (self.name, self.dvalue)
else:
coords[self.dname] = self.dvalue
return coords

@property
def ndim(self) -> int:
Expand All @@ -104,6 +106,14 @@ def size(self) -> int:
def length(self) -> FloatArray:
return np.squeeze(abs(np.diff(self.bounds, axis=1)))

@property
def directional_bounds(self):
# Only flip bounds if needed
if self.flipped:
return self.bounds[::-1, :].copy()
else:
return self.bounds

def flip_if_needed(self, index: IntArray) -> IntArray:
if self.flipped:
return self.size - index - 1
Expand Down Expand Up @@ -462,8 +472,8 @@ def convert_to(self, matched_type: Any) -> Any:
return self
elif matched_type == UnstructuredGrid2d:
ugrid2d = Ugrid2d.from_structured_bounds(
self.xbounds.bounds,
self.ybounds.bounds,
self.xbounds.directional_bounds,
self.ybounds.directional_bounds,
)
return UnstructuredGrid2d(ugrid2d)
else:
Expand Down

0 comments on commit d242307

Please sign in to comment.