Merge pull request #120 from Deltares/Issue_99

Issue 99

tests/conftest.py

@@ -19,6 +19,11 @@
     grid_data_c_2d,
     grid_data_d,
     grid_data_d_1d,
+    grid_data_dask_expected,
+    grid_data_dask_expected_layered,
+    grid_data_dask_source,
+    grid_data_dask_source_layered,
+    grid_data_dask_target,
     grid_data_e,
     grid_data_e_1d,
     quads_0_25,

tests/fixtures/fixture_regridder.py

@@ -1,3 +1,4 @@
+import dask.array as da
 import numpy as np
 import pytest
 import xarray as xr
@@ -337,3 +338,74 @@ def expected_results_linear():
             "dy": -50.0,
         },
     )
+
+
+@pytest.fixture(scope="function")
+def grid_data_dask_source():
+    data = np.arange(10000).reshape((100, 100))
+    data = da.from_array(data, chunks=(10, (10, 30, 60)))
+    return xr.DataArray(
+        data=data,
+        dims=["y", "x"],
+        coords={
+            "y": -np.arange(100),
+            "x": np.arange(100),
+            "dx": 1.0,
+            "dy": -1.0,
+        },
+    )
+
+
+@pytest.fixture(scope="function")
+def grid_data_dask_source_layered():
+    data = np.arange(30000).reshape((3, 100, 100))
+    data = da.from_array(data, chunks=(3, 10, (10, 30, 60)))
+    return xr.DataArray(
+        data=data,
+        dims=["layer", "y", "x"],
+        coords={
+            "layer": np.arange(3),
+            "y": -np.arange(100),
+            "x": np.arange(100),
+            "dx": 1.0,
+            "dy": -1.0,
+        },
+    )
+
+
+@pytest.fixture(scope="function")
+def grid_data_dask_target():
+    data = np.zeros(100).reshape((10, 10))
+    data = da.from_array(data, chunks=(10, (5, 5)))
+    return xr.DataArray(
+        data=data,
+        dims=["y", "x"],
+        coords={
+            "y": -np.arange(10) * 10,
+            "x": np.arange(10) * 10,
+            "dx": 10.0,
+            "dy": -10.0,
+        },
+    )
+
+
+@pytest.fixture(scope="function")
+def grid_data_dask_expected():
+    data1 = np.tile(np.arange(0.0, 100.0, 10.0), reps=10).reshape((10, 10))
+    data2 = np.repeat(np.arange(0.0, 10000.0, 1000.0), repeats=10).reshape((10, 10))
+    data = data1 + data2
+    return xr.DataArray(
+        data=data,
+        dims=["y", "x"],
+        coords={
+            "y": -np.arange(10) * 10,
+            "x": np.arange(10) * 10,
+            "dx": 10.0,
+            "dy": -10.0,
+        },
+    )
+
+
+@pytest.fixture(scope="function")
+def grid_data_dask_expected_layered(grid_data_dask_expected):
+    return grid_data_dask_expected.expand_dims(dim={"layer": np.arange(3)})

tests/test_regrid/test_regridder.py

@@ -192,3 +192,24 @@ def test_regridder_from_dataset(cls, disk, quads_1):
     new_regridder = cls.from_dataset(dataset)
     new_result = new_regridder.regrid(disk)
     assert np.array_equal(new_result.values, result.values, equal_nan=True)
+
+
+def test_regridder_daks_arrays(
+    grid_data_dask_source,
+    grid_data_dask_source_layered,
+    grid_data_dask_target,
+    grid_data_dask_expected,
+    grid_data_dask_expected_layered,
+):
+    regridder = CentroidLocatorRegridder(
+        source=grid_data_dask_source, target=grid_data_dask_target
+    )
+    result = regridder.regrid(grid_data_dask_source)
+    assert result.equals(grid_data_dask_expected)
+
+    # with broadcasting
+    regridder = CentroidLocatorRegridder(
+        source=grid_data_dask_source_layered, target=grid_data_dask_target
+    )
+    result = regridder.regrid(grid_data_dask_source_layered)
+    assert result.isel(layer=0).equals(grid_data_dask_expected_layered.isel(layer=0))

xugrid/regrid/regridder.py

@@ -144,8 +144,17 @@ def _regrid_array(self, source):
         source = source.reshape((-1, source_grid.size))
 
         size = self._target.size
+
         if isinstance(source, DaskArray):
-            chunks = source.chunks[: -source_grid.ndim] + (self._target.shape,)
+            # It's possible that the topology dimensions are chunked (e.g. from
+            # reading multiple partitions). The regrid operation does not
+            # support this, since we might need multiple source chunks for a
+            # single target chunk, which destroys the 1:1 relation between
+            # chunks. Here we ensure that the topology dimensions are contained
+            # in a single contiguous chunk.
+            contiguous_chunks = (source.chunks[0], (source.shape[-1],))
+            source = source.rechunk(contiguous_chunks)
+            chunks = source.chunks[:-1] + (self._target.size,)
             out = dask.array.map_blocks(
                 self._regrid,  # func
                 source,  # *args
@@ -162,7 +171,6 @@ def _regrid_array(self, source):
                 "Expected dask.array.Array or numpy.ndarray. Received: "
                 f"{type(source).__name__}"
             )
-
         # E.g.: sizes of ("time", "layer") + ("y", "x")
         out_shape = first_dims_shape + self._target.shape
         return out.reshape(out_shape)
@@ -203,12 +211,7 @@ def regrid(self, object) -> UgridDataArray:
         if type(self._target) is StructuredGrid2d:
             source_dims = ("y", "x")
             regridded = self.regrid_dataarray(object, source_dims)
-            regridded = regridded.assign_coords(
-                coords={
-                    "y": np.flip(self._target.ybounds.midpoints),
-                    "x": self._target.xbounds.midpoints,
-                }
-            )
+            regridded = regridded.assign_coords(coords=self._target.coords)
             return regridded
         else:
             source_dims = (object.ugrid.grid.face_dimension,)

xugrid/regrid/structured.py

@@ -74,7 +74,16 @@ def __init__(self, obj: Union[xr.DataArray, xr.Dataset], name: str):
         self.bounds = bounds
         self.flipped = flipped
         self.side = side
-        self.grid = obj
+        self.dname = size_name
+        self.dvalue = obj[size_name].values
+        self.index = obj.indexes[name].values
+
+    @property
+    def coords(self) -> dict:
+        return {
+            self.name: self.index,
+            self.dname: self.dvalue,
+        }
 
     @property
     def ndim(self) -> int:
@@ -420,6 +429,10 @@ def __init__(
         self.xbounds = StructuredGrid1d(obj, name_x)
         self.ybounds = StructuredGrid1d(obj, name_y)
 
+    @property
+    def coords(self) -> dict:
+        return self.ybounds.coords | self.xbounds.coords
+
     @property
     def ndim(self) -> int:
         return 2