Implement Grid3D.find_index() and enable da.sel() for dfs3 data

src/mikeio/dataset/_dataarray.py

@@ -841,15 +841,21 @@ def sel(
         if len(kwargs) > 0:
             idx = self.geometry.find_index(**kwargs)
 
-            # TODO this seems fragile
             if isinstance(idx, tuple):
-                # TODO: support for dfs3
-                assert len(idx) == 2
-                ii, jj = idx
-                if jj is not None:
-                    da = da.isel(y=jj)
-                if ii is not None:
-                    da = da.isel(x=ii)
+                if len(idx) == 3:
+                    ii, jj, kk = idx
+                    if kk is not None:
+                        da = da.isel(z=kk)
+                    if jj is not None:
+                        da = da.isel(y=jj)
+                    if ii is not None:
+                        da = da.isel(x=ii)
+                elif len(idx) == 2:
+                    ii, jj = idx
+                    if jj is not None:
+                        da = da.isel(y=jj)
+                    if ii is not None:
+                        da = da.isel(x=ii)
             else:
                 da = da.isel(idx, axis="space")
 
@@ -885,6 +891,8 @@ def _sel_with_slice(self, kwargs: Mapping[str, slice]) -> DataArray:
                         pos = 0
                     if k == "y":
                         pos = 1
+                    if k == "z":
+                        pos = 2
 
                 start = idx_start[pos][0] if idx_start is not None else None
                 stop = idx_stop[pos][0] if idx_stop is not None else None

src/mikeio/spatial/_grid_geometry.py

@@ -1223,16 +1223,152 @@ def orientation(self) -> float:
         """Grid orientation."""
         return self._orientation
 
+    @property
+    def bbox(self) -> BoundingBox:
+        """Bounding box (left, bottom, right, top).
+
+        Note: not the same as the cell center values (x0,y0,x1,y1)!
+        """
+        if self._is_rotated:
+            raise NotImplementedError("Only available if orientation = 0")
+        left = self.x[0] - self.dx / 2
+        bottom = self.y[0] - self.dy / 2
+        right = self.x[-1] + self.dx / 2
+        top = self.y[-1] + self.dy / 2
+        return BoundingBox(left, bottom, right, top)
+
+    def contains(self, coords: ArrayLike) -> Any:
+        """Test if a list of xy-points are inside grid.
+
+        Parameters
+        ----------
+        coords : array(float)
+            xy-coordinate of points given as n-by-2 array
+
+        Returns
+        -------
+        bool array
+            True for points inside, False otherwise
+
+        """
+        coords = np.atleast_2d(coords)
+        x = coords[:, 0]
+        y = coords[:, 1]
+        xinside = (self.bbox.left <= x) & (x <= self.bbox.right)
+        yinside = (self.bbox.bottom <= y) & (y <= self.bbox.top)
+        return xinside & yinside
+
+    def __contains__(self, pt: Any) -> Any:
+        return self.contains(pt)
+
+    def _xy_to_index(self, xy: ArrayLike) -> tuple[np.ndarray, np.ndarray]:
+        """Find nearest (i, j) indices for xy points."""
+        xy = np.atleast_2d(xy)
+        x = xy[:, 0]
+        y = xy[:, 1]
+
+        inside = self.contains(xy)
+        if np.any(~inside):
+            raise OutsideModelDomainError(x=x[~inside], y=y[~inside])
+
+        ii = np.floor((x - (self.x[0] - self.dx / 2)) / self.dx).astype(int)
+        jj = np.floor((y - (self.y[0] - self.dy / 2)) / self.dy).astype(int)
+        return ii, jj
+
+    def _bbox_to_index(self, bbox: BoundingBox) -> tuple[range, range]:
+        """Find subarea within this geometry."""
+        if not bbox.overlaps(self.bbox):
+            raise ValueError("area is outside grid")
+
+        mask = (self.x >= bbox.left) & (self.x <= bbox.right)
+        ii = np.where(mask)[0]
+        mask = (self.y >= bbox.bottom) & (self.y <= bbox.top)
+        jj = np.where(mask)[0]
+
+        i = range(ii[0], ii[-1] + 1)
+        j = range(jj[0], jj[-1] + 1)
+        return i, j
+
     def find_index(
-        self, coords: Any = None, layers: Any = None, area: Any = None
-    ) -> Any:
+        self,
+        x: float | None = None,
+        y: float | None = None,
+        z: float | None = None,
+        coords: ArrayLike | None = None,
+        area: tuple[float, float, float, float] | None = None,
+        layers: Any = None,
+    ) -> tuple[Any, Any, Any]:
+        """Find nearest index (i, j, k) of point(s).
+
+        Parameters
+        ----------
+        x : float, optional
+            x-coordinate of point
+        y : float, optional
+            y-coordinate of point
+        z : float, optional
+            z-coordinate of point
+        coords : array(float), optional
+            xy-coordinates of points given as n-by-2 array
+        area : tuple(float), optional
+            xy-coordinates of bounding box (left, bottom, right, top)
+        layers : Any, optional
+            Not yet implemented
+
+        Returns
+        -------
+        tuple
+            (i, j, k) indices — any component may be None if not specified
+
+        """
         if layers is not None:
             raise NotImplementedError(
                 f"Layer slicing is not yet implemented. Use the mikeio.read('file.dfs3', layers='{layers}')"
             )
-        raise NotImplementedError(
-            "Not yet implemented for Grid3D. Please use mikeio.read('file.dfs3') and its arguments instead."
-        )
+
+        if x is not None and not np.isscalar(x):
+            raise ValueError(
+                f"{x=} is not a scalar value, use the coords argument instead"
+            )
+        if y is not None and not np.isscalar(y):
+            raise ValueError(
+                f"{y=} is not a scalar value, use the coords argument instead"
+            )
+
+        kk: Any = None
+        if z is not None:
+            if z < self.z[0] or z > self.z[-1]:
+                raise OutsideModelDomainError(x=None, y=None, z=z)
+            kk = np.atleast_1d(np.argmin(np.abs(self.z - z)))
+
+        if x is not None and y is not None:
+            if coords is not None:
+                raise ValueError("x,y and coords cannot be given at the same time!")
+            coords = np.column_stack([np.atleast_1d(x), np.atleast_1d(y)])
+
+        if x is not None and y is None and coords is None:
+            if x < self.x[0] or x > self.x[-1]:
+                raise OutsideModelDomainError(x=x, y=None)
+            return np.atleast_1d(np.argmin(np.abs(self.x - x))), None, kk
+
+        if y is not None and x is None and coords is None:
+            if y < self.y[0] or y > self.y[-1]:
+                raise OutsideModelDomainError(x=None, y=y)
+            return None, np.atleast_1d(np.argmin(np.abs(self.y - y))), kk
+
+        if coords is not None:
+            ii, jj = self._xy_to_index(coords)
+            return ii, jj, kk
+
+        if area is not None:
+            bbox = BoundingBox.parse(area)
+            i, j = self._bbox_to_index(bbox)
+            return i, j, kk
+
+        if z is not None:
+            return None, None, kk
+
+        raise ValueError("Provide x, y, z, coords, or area")
 
     def isel(
         self, idx: int | np.ndarray, axis: int

tests/test_dataarray.py

@@ -472,6 +472,48 @@ def test_dataarray_grid3d_indexing() -> None:
 
     # TODO: wait for merge of https://github.com/DHI/mikeio/pull/311
     # assert isinstance(da[:, 1, ::3, :].geometry, mikeio.Grid2D)
+
+
+def test_da_sel_xyz_grid3d() -> None:
+    da = mikeio.read("tests/testdata/test_dfs3.dfs3")[0]
+    x0, y0 = da.geometry.x[2], da.geometry.y[3]
+    z0 = da.geometry.z[1]
+    da2 = da.sel(x=x0, y=y0, z=z0)
+    assert da2.shape == (da.n_timesteps,)
+
+
+def test_da_sel_x_grid3d() -> None:
+    da = mikeio.read("tests/testdata/test_dfs3.dfs3")[0]
+    x0 = da.geometry.x[2]
+    da2 = da.sel(x=x0)
+    # selecting x reduces x dim, keeps z and y
+    assert da2.shape == (da.n_timesteps, da.geometry.nz, da.geometry.ny)
+
+
+def test_da_sel_z_grid3d() -> None:
+    da = mikeio.read("tests/testdata/test_dfs3.dfs3")[0]
+    z0 = da.geometry.z[1]
+    da2 = da.sel(z=z0)
+    # selecting z reduces z dim, keeps x and y => Grid2D
+    assert isinstance(da2.geometry, mikeio.Grid2D)
+    assert da2.shape == (da.n_timesteps, da.geometry.ny, da.geometry.nx)
+
+
+def test_da_sel_area_grid3d() -> None:
+    da = mikeio.read("tests/testdata/test_dfs3.dfs3")[0]
+    x = da.geometry.x
+    y = da.geometry.y
+    area = (x[1], y[1], x[4], y[4])
+    da2 = da.sel(area=area)
+    assert isinstance(da2.geometry, mikeio.Grid3D)
+
+
+def test_da_sel_slice_grid3d() -> None:
+    da = mikeio.read("tests/testdata/test_dfs3.dfs3")[0]
+    x = da.geometry.x
+    z = da.geometry.z
+    da2 = da.sel(x=slice(x[1], x[4]), z=slice(z[0], z[2]))
+    assert isinstance(da2.geometry, mikeio.Grid3D)
     # assert isinstance(da[:, 1, -3, 4:].geometry, mikeio.Grid2D)
 
 

tests/test_geometry_grid.py

@@ -419,3 +419,87 @@ def test_bad_projection_raises_error() -> None:
 def test_grid3d_repr() -> None:
     g = Grid3D(nx=2, ny=2, nz=2, dx=1, dy=1, dz=1, projection="UTM-33")
     assert "Grid3D" in repr(g)
+
+
+def test_grid3d_bbox() -> None:
+    g = Grid3D(nx=5, ny=10, nz=3, dx=100.0, dy=100.0, dz=1.0)
+    bb = g.bbox
+    assert bb.left == 0.0
+    assert bb.bottom == 0.0
+    assert bb.right == 500.0
+    assert bb.top == 1000.0
+
+
+def test_grid3d_contains() -> None:
+    g = Grid3D(nx=5, ny=10, nz=3, dx=100.0, dy=100.0, dz=1.0)
+    assert g.contains([250, 500])
+    assert not g.contains([600, 500])
+
+
+def test_grid3d_find_index_x_only() -> None:
+    g = Grid3D(nx=5, ny=10, nz=3, dx=100.0, dy=100.0, dz=1.0)
+    ii, jj, kk = g.find_index(x=150.0)
+    assert ii == 1
+    assert jj is None
+    assert kk is None
+
+
+def test_grid3d_find_index_y_only() -> None:
+    g = Grid3D(nx=5, ny=10, nz=3, dx=100.0, dy=100.0, dz=1.0)
+    ii, jj, kk = g.find_index(y=350.0)
+    assert ii is None
+    assert jj == 3
+    assert kk is None
+
+
+def test_grid3d_find_index_z_only() -> None:
+    g = Grid3D(nx=5, ny=10, nz=3, dx=100.0, dy=100.0, dz=1.0)
+    ii, jj, kk = g.find_index(z=1.5)
+    assert ii is None
+    assert jj is None
+    assert kk == 1
+
+
+def test_grid3d_find_index_xy() -> None:
+    g = Grid3D(nx=5, ny=10, nz=3, dx=100.0, dy=100.0, dz=1.0)
+    ii, jj, kk = g.find_index(x=150.0, y=350.0)
+    assert ii == 1
+    assert jj == 3
+    assert kk is None
+
+
+def test_grid3d_find_index_xyz() -> None:
+    g = Grid3D(nx=5, ny=10, nz=3, dx=100.0, dy=100.0, dz=1.0)
+    ii, jj, kk = g.find_index(x=150.0, y=350.0, z=1.5)
+    assert ii == 1
+    assert jj == 3
+    assert kk == 1
+
+
+def test_grid3d_find_index_coords() -> None:
+    g = Grid3D(nx=5, ny=10, nz=3, dx=100.0, dy=100.0, dz=1.0)
+    xy = np.array([[150.0, 350.0], [250.0, 550.0]])
+    ii, jj, kk = g.find_index(coords=xy)
+    assert ii[0] == 1
+    assert jj[0] == 3
+    assert ii[1] == 2
+    assert jj[1] == 5
+    assert kk is None
+
+
+def test_grid3d_find_index_area() -> None:
+    g = Grid3D(nx=5, ny=10, nz=3, dx=100.0, dy=100.0, dz=1.0)
+    ii, jj, kk = g.find_index(area=(100, 200, 300, 600))
+    assert isinstance(ii, range)
+    assert isinstance(jj, range)
+    assert kk is None
+
+
+def test_grid3d_find_index_outside() -> None:
+    g = Grid3D(nx=5, ny=10, nz=3, dx=100.0, dy=100.0, dz=1.0)
+    with pytest.raises(OutsideModelDomainError):
+        g.find_index(x=-100.0)
+    with pytest.raises(OutsideModelDomainError):
+        g.find_index(y=2000.0)
+    with pytest.raises(OutsideModelDomainError):
+        g.find_index(z=100.0)