feat: add xarray grid coordinate indexing (#207)

- Added x, y, z world coordinates to StructuredGrid
- Coordinates computed from grid geometry (delr, delc, top, botm)
- Enables coordinate-based selection: grid.botm.sel(x=100, method='nearest')
- Added comprehensive tests for coordinate indexing
- All grid data arrays now have spatial coordinates attached

Author: claude

flopy4/mf6/utils/grid.py

@@ -23,12 +23,21 @@ def __init__(self, *args, **kwargs):
             "ncol": "j",
             "nodes": "node",
         }
+
+        # Compute world coordinates (x, y, z)
+        world_coords = self._compute_world_coordinates()
+
+        # Index coordinates
         self._coords = {
             "k": xr.DataArray(np.arange(self.nlay, dtype=int), dims=("nlay",)),
             "i": xr.DataArray(np.arange(self.nrow, dtype=int), dims=("nrow",)),
             "j": xr.DataArray(np.arange(self.ncol, dtype=int), dims=("ncol",)),
             "node": xr.DataArray(np.arange(self.nnodes, dtype=int), dims=("nodes",)),
         }
+
+        # Add world coordinates
+        self._coords.update(world_coords)
+
         data_vars = {
             "delr": self.delr,
             "delc": self.delc,
@@ -43,8 +52,56 @@ def __init__(self, *args, **kwargs):
             .set_xindex("i", PandasIndex)
             .set_xindex("j", PandasIndex)
             .set_xindex("node", PandasIndex)
+            .set_xindex("x", PandasIndex)
+            .set_xindex("y", PandasIndex)
+            .set_xindex("z", PandasIndex)
         )
 
+    def _compute_world_coordinates(self) -> dict:
+        """
+        Compute x, y, z world coordinates from grid geometry.
+
+        Returns
+        -------
+        dict
+            Dictionary with 'x', 'y', 'z' coordinate DataArrays
+        """
+        # Get grid extent
+        xmin, xmax, ymin, ymax = self.extent
+
+        # Compute x coordinates (cell centers)
+        delr = np.atleast_1d(self.delr[0] if hasattr(self.delr, '__getitem__') and not isinstance(self.delr, np.ndarray) else self.delr)
+        if delr.size == 1:
+            delr = np.full(self.ncol, delr[0])
+        x = xmin + np.cumsum(delr) - 0.5 * delr
+
+        # Compute y coordinates (cell centers)
+        delc = np.atleast_1d(self.delc[0] if hasattr(self.delc, '__getitem__') and not isinstance(self.delc, np.ndarray) else self.delc)
+        if delc.size == 1:
+            delc = np.full(self.nrow, delc[0])
+        y = ymax - np.cumsum(delc) + 0.5 * delc
+
+        # Compute z coordinates (layer centers)
+        # Use top and botm to compute layer centers
+        top_2d = np.atleast_2d(self.top)
+        botm_3d = np.atleast_3d(self.botm).reshape(self.nlay, self.nrow, self.ncol)
+
+        # Layer center z coordinates: average of top and bottom per layer
+        z = np.zeros(self.nlay)
+        for k in range(self.nlay):
+            if k == 0:
+                layer_top = top_2d.mean()
+            else:
+                layer_top = botm_3d[k - 1].mean()
+            layer_bot = botm_3d[k].mean()
+            z[k] = (layer_top + layer_bot) / 2.0
+
+        return {
+            "x": xr.DataArray(x, dims=("ncol",)),
+            "y": xr.DataArray(y, dims=("nrow",)),
+            "z": xr.DataArray(z, dims=("nlay",)),
+        }
+
     @property
     def dataset(self) -> xr.Dataset:
         return self._dataset
@@ -55,21 +112,21 @@ def delc(self):
             return None
         dims = ("ncol",)
         coord_name = self._dims_coords[dims[0]]
-        coords = coords = {coord_name: self._coords[coord_name]}
+        coords = {coord_name: self._coords[coord_name], "x": self._coords["x"]}
         return xr.DataArray(super().delc, coords=coords, dims=dims).set_xindex(
             coord_name, PandasIndex
-        )
+        ).set_xindex("x", PandasIndex)
 
     @property
     def delr(self):
         if self.__delr is None:
             return None
         dims = ("nrow",)
         coord_name = self._dims_coords[dims[0]]
-        coords = {coord_name: self._coords[coord_name]}
+        coords = {coord_name: self._coords[coord_name], "y": self._coords["y"]}
         return xr.DataArray(super().delr, coords=coords, dims=dims).set_xindex(
             coord_name, PandasIndex
-        )
+        ).set_xindex("y", PandasIndex)
 
     @property
     def delz(self):
@@ -80,22 +137,29 @@ def delz(self):
             self._dims_coords[dims[2]],
         )
         coords = {coord_name: self._coords[coord_name] for coord_name in coord_names}
+        coords.update({"x": self._coords["x"], "y": self._coords["y"], "z": self._coords["z"]})
         return (
             xr.DataArray(super().delz, coords=coords, dims=dims)
             .set_xindex(coord_names[0], PandasIndex)
             .set_xindex(coord_names[1], PandasIndex)
             .set_xindex(coord_names[2], PandasIndex)
+            .set_xindex("x", PandasIndex)
+            .set_xindex("y", PandasIndex)
+            .set_xindex("z", PandasIndex)
         )
 
     @property
     def top(self):
         dims = ("nrow", "ncol")
         coord_names = (self._dims_coords[dims[0]], self._dims_coords[dims[1]])
         coords = {coord_name: self._coords[coord_name] for coord_name in coord_names}
+        coords.update({"x": self._coords["x"], "y": self._coords["y"]})
         return (
             xr.DataArray(super().top, coords=coords, dims=dims)
             .set_xindex(coord_names[0], PandasIndex)
             .set_xindex(coord_names[1], PandasIndex)
+            .set_xindex("x", PandasIndex)
+            .set_xindex("y", PandasIndex)
         )
 
     @property
@@ -107,11 +171,15 @@ def botm(self):
             self._dims_coords[dims[2]],
         )
         coords = {coord_name: self._coords[coord_name] for coord_name in coord_names}
+        coords.update({"x": self._coords["x"], "y": self._coords["y"], "z": self._coords["z"]})
         return (
             xr.DataArray(super().botm, coords=coords, dims=dims)
             .set_xindex(coord_names[0], PandasIndex)
             .set_xindex(coord_names[1], PandasIndex)
             .set_xindex(coord_names[2], PandasIndex)
+            .set_xindex("x", PandasIndex)
+            .set_xindex("y", PandasIndex)
+            .set_xindex("z", PandasIndex)
         )
 
     @property
@@ -123,11 +191,15 @@ def idomain(self):
             self._dims_coords[dims[2]],
         )
         coords = {coord_name: self._coords[coord_name] for coord_name in coord_names}
+        coords.update({"x": self._coords["x"], "y": self._coords["y"], "z": self._coords["z"]})
         return (
             xr.DataArray(super().idomain, coords=coords, dims=dims)
             .set_xindex(coord_names[0], PandasIndex)
             .set_xindex(coord_names[1], PandasIndex)
             .set_xindex(coord_names[2], PandasIndex)
+            .set_xindex("x", PandasIndex)
+            .set_xindex("y", PandasIndex)
+            .set_xindex("z", PandasIndex)
         )
 
 

test/test_mf6_component.py

@@ -627,3 +627,96 @@ def test_to_xarray_on_context(function_tmpdir):
     assert np.array_equal(dt.per, [0])
     assert dt.attrs["filename"] == "mfsim.nam"
     assert dt.attrs["workspace"] == Path(function_tmpdir)
+
+
+def test_grid_coordinate_indexing():
+    """Test that grid data can be indexed using x, y, z world coordinates."""
+    # Create a simple structured grid
+    grid = StructuredGrid(
+        nlay=3,
+        nrow=10,
+        ncol=10,
+        delr=100.0,  # 100 m cell width
+        delc=100.0,  # 100 m cell height
+        top=100.0,
+        botm=[0.0, -50.0, -100.0],
+    )
+
+    # Test that world coordinates are present
+    assert "x" in grid.dataset.coords
+    assert "y" in grid.dataset.coords
+    assert "z" in grid.dataset.coords
+
+    # Test coordinate shapes
+    assert len(grid.dataset.coords["x"]) == 10  # ncol
+    assert len(grid.dataset.coords["y"]) == 10  # nrow
+    assert len(grid.dataset.coords["z"]) == 3   # nlay
+
+    # Test coordinate values
+    # X coordinates should be cell centers: 50, 150, 250, ..., 950
+    expected_x = np.array([50.0, 150.0, 250.0, 350.0, 450.0, 550.0, 650.0, 750.0, 850.0, 950.0])
+    np.testing.assert_allclose(grid.dataset.coords["x"].values, expected_x)
+
+    # Y coordinates should be cell centers from top: 950, 850, 750, ..., 50
+    expected_y = np.array([950.0, 850.0, 750.0, 650.0, 550.0, 450.0, 350.0, 250.0, 150.0, 50.0])
+    np.testing.assert_allclose(grid.dataset.coords["y"].values, expected_y)
+
+    # Z coordinates should be layer centers
+    # Layer 0: (100 + 0) / 2 = 50
+    # Layer 1: (0 + (-50)) / 2 = -25
+    # Layer 2: (-50 + (-100)) / 2 = -75
+    expected_z = np.array([50.0, -25.0, -75.0])
+    np.testing.assert_allclose(grid.dataset.coords["z"].values, expected_z)
+
+    # Test coordinate-based selection using nearest
+    # Select near x=250 (should get col=2, which has x=250)
+    botm_at_x250 = grid.botm.sel(x=250.0, method="nearest")
+    assert botm_at_x250.shape == (3, 10)  # (nlay, nrow)
+
+    # Select near y=650 (should get row=3, which has y=650)
+    botm_at_y650 = grid.botm.sel(y=650.0, method="nearest")
+    assert botm_at_y650.shape == (3, 10)  # (nlay, ncol)
+
+    # Select near z=50 (should get lay=0, which has z=50)
+    botm_at_z50 = grid.botm.sel(z=50.0, method="nearest")
+    assert botm_at_z50.shape == (10, 10)  # (nrow, ncol)
+
+    # Test combined selection
+    botm_at_point = grid.botm.sel(x=250.0, y=650.0, z=50.0, method="nearest")
+    assert botm_at_point.shape == ()  # scalar
+
+    # Verify the value makes sense
+    # This should be layer 0, row 3, col 2 -> botm[0]
+    expected_value = 0.0  # botm[0] = 0.0
+    assert float(botm_at_point) == expected_value
+
+
+def test_grid_coordinate_indexing_in_dis():
+    """Test that Dis package data also has coordinate indexing."""
+    time = Time(perlen=[1.0], nstp=[1])
+    dis = Dis(
+        nlay=2,
+        nrow=5,
+        ncol=5,
+        delr=10.0,
+        delc=10.0,
+        top=10.0,
+        botm=[0.0, -10.0],
+    )
+
+    # Convert to grid to access coordinates
+    grid = dis.to_grid()
+
+    # Test that coordinates are available
+    assert "x" in grid.dataset.coords
+    assert "y" in grid.dataset.coords
+    assert "z" in grid.dataset.coords
+
+    # Test coordinate-based selection on botm
+    # X coordinates: 5, 15, 25, 35, 45
+    botm_near_x15 = grid.botm.sel(x=15.0, method="nearest")
+    assert botm_near_x15.shape == (2, 5)  # (nlay, nrow)
+
+    # Y coordinates: 45, 35, 25, 15, 5
+    botm_near_y25 = grid.botm.sel(y=25.0, method="nearest")
+    assert botm_near_y25.shape == (2, 5)  # (nlay, ncol)