feat(grid): add array support to .intersect() (#2646)

Add vectorized array support to the intersect() method on all grid types

Author: aacovski

autotest/test_grid.py

@@ -11,6 +11,7 @@
 from matplotlib import pyplot as plt
 from modflow_devtools.markers import requires_exe, requires_pkg
 from modflow_devtools.misc import has_pkg
+from scipy.spatial import Delaunay
 
 from autotest.test_dis_cases import case_dis, case_disv
 from autotest.test_grid_cases import GridCases
@@ -77,6 +78,17 @@ def minimal_vertex_grid_info(minimal_unstructured_grid_info):
     return d
 
 
+@pytest.fixture
+def simple_structured_grid():
+    """Create a simple 10x10 structured grid for testing."""
+    nrow, ncol = 10, 10
+    delr = np.ones(ncol) * 10.0
+    delc = np.ones(nrow) * 10.0
+    top = np.ones((nrow, ncol)) * 10.0
+    botm = np.zeros((1, nrow, ncol))
+    return StructuredGrid(delr=delr, delc=delc, top=top, botm=botm)
+
+
 def test_rotation():
     m = Modflow(rotation=20.0)
     dis = ModflowDis(
@@ -419,6 +431,127 @@ def test_unstructured_xyz_intersect(example_data_path):
             raise AssertionError("Unstructured grid intersection failed")
 
 
+def test_structured_grid_intersect_array(simple_structured_grid):
+    """Test StructuredGrid.intersect() with array inputs."""
+    grid = simple_structured_grid
+
+    # Test array input
+    x = np.array([25.0, 50.0, 75.0])
+    y = np.array([25.0, 50.0, 75.0])
+    rows, cols = grid.intersect(x, y, forgive=True)
+
+    # Verify array output
+    assert isinstance(rows, np.ndarray)
+    assert isinstance(cols, np.ndarray)
+    assert len(rows) == 3
+    assert len(cols) == 3
+
+    # Verify equivalence with scalar calls
+    for i in range(len(x)):
+        row_scalar, col_scalar = grid.intersect(x[i], y[i], forgive=True)
+        assert rows[i] == row_scalar
+        assert cols[i] == col_scalar
+
+    # Test out-of-bounds with forgive
+    x_mixed = np.array([50.0, 150.0])
+    y_mixed = np.array([50.0, 50.0])
+    rows_mixed, cols_mixed = grid.intersect(x_mixed, y_mixed, forgive=True)
+    assert not np.isnan(rows_mixed[0])  # First point valid
+    assert np.isnan(rows_mixed[1])  # Second point out of bounds
+
+
+def test_vertex_grid_intersect_array():
+    """Test VertexGrid.intersect() with array inputs."""
+    # Create a simple vertex grid using Delaunay triangulation
+    np.random.seed(42)
+    n_points = 50
+    x_verts = np.random.uniform(0, 100, n_points)
+    y_verts = np.random.uniform(0, 100, n_points)
+    points = np.column_stack([x_verts, y_verts])
+
+    tri = Delaunay(points)
+    vertices = [[i, x_verts[i], y_verts[i]] for i in range(len(x_verts))]
+    cell2d = [[i] + list(tri.simplices[i]) for i in range(len(tri.simplices))]
+
+    ncells = len(cell2d)
+    top = np.ones(ncells) * 10.0
+    botm = np.zeros(ncells)
+    grid = VertexGrid(vertices=vertices, cell2d=cell2d, top=top, botm=botm)
+
+    # Test array input
+    x = np.array([25.0, 50.0, 75.0])
+    y = np.array([25.0, 50.0, 75.0])
+    results = grid.intersect(x, y, forgive=True)
+
+    # Verify array output
+    assert isinstance(results, np.ndarray)
+    assert len(results) == 3
+
+    # Verify equivalence with scalar calls
+    for i in range(len(x)):
+        result_scalar = grid.intersect(x[i], y[i], forgive=True)
+        if np.isnan(results[i]) and np.isnan(result_scalar):
+            continue
+        assert results[i] == result_scalar
+
+    # Test out-of-bounds with forgive
+    x_mixed = np.array([50.0, 200.0])
+    y_mixed = np.array([50.0, 50.0])
+    results_mixed = grid.intersect(x_mixed, y_mixed, forgive=True)
+    assert np.isnan(results_mixed[1])  # Second point out of bounds
+
+
+def test_unstructured_grid_intersect_array():
+    """Test UnstructuredGrid.intersect() with array inputs."""
+    # Create a simple unstructured grid using Delaunay triangulation
+    np.random.seed(42)
+    n_points = 50
+    x_verts = np.random.uniform(0, 100, n_points)
+    y_verts = np.random.uniform(0, 100, n_points)
+    points = np.column_stack([x_verts, y_verts])
+
+    tri = Delaunay(points)
+    vertices = [[i, x_verts[i], y_verts[i]] for i in range(len(x_verts))]
+    iverts = tri.simplices.tolist()
+
+    xcenters = np.mean(points[tri.simplices], axis=1)[:, 0]
+    ycenters = np.mean(points[tri.simplices], axis=1)[:, 1]
+
+    ncells = len(iverts)
+    top = np.ones(ncells) * 10.0
+    botm = np.zeros(ncells)
+    grid = UnstructuredGrid(
+        vertices=vertices,
+        iverts=iverts,
+        xcenters=xcenters,
+        ycenters=ycenters,
+        top=top,
+        botm=botm,
+    )
+
+    # Test array input
+    x = np.array([25.0, 50.0, 75.0])
+    y = np.array([25.0, 50.0, 75.0])
+    results = grid.intersect(x, y, forgive=True)
+
+    # Verify array output
+    assert isinstance(results, np.ndarray)
+    assert len(results) == 3
+
+    # Verify equivalence with scalar calls
+    for i in range(len(x)):
+        result_scalar = grid.intersect(x[i], y[i], forgive=True)
+        if np.isnan(results[i]) and np.isnan(result_scalar):
+            continue
+        assert results[i] == result_scalar
+
+    # Test out-of-bounds with forgive
+    x_mixed = np.array([50.0, 200.0])
+    y_mixed = np.array([50.0, 50.0])
+    results_mixed = grid.intersect(x_mixed, y_mixed, forgive=True)
+    assert np.isnan(results_mixed[1])  # Second point out of bounds
+
+
 @pytest.mark.parametrize("spc_file", ["grd.spc", "grdrot.spc"])
 def test_structured_from_gridspec(example_data_path, spc_file):
     fn = example_data_path / "specfile" / spc_file
@@ -1522,3 +1655,20 @@ def test_unstructured_iverts_cleanup():
 
     if clean_ugrid.nvert != cleaned_vert_num:
         raise AssertionError("Improper number of vertices for cleaned 'shared' iverts")
+
+
+def test_structured_grid_intersect_edge_case(simple_structured_grid):
+    """Test StructuredGrid.intersect() edge case: out-of-bounds x,y with z.
+
+    This tests the specific case where x,y are out of bounds and z is provided.
+    The expected behavior is to return (None, nan, nan).
+    """
+    grid = simple_structured_grid
+
+    # Test out-of-bounds x,y with z coordinate
+    lay, row, col = grid.intersect(-50.0, -50.0, z=5.0, forgive=True)
+
+    # Expected behavior: lay=None, row=nan, col=nan
+    assert lay is None, f"Expected lay=None, got {lay}"
+    assert np.isnan(row), f"Expected row=nan, got {row}"
+    assert np.isnan(col), f"Expected col=nan, got {col}"

flopy/discretization/structuredgrid.py

@@ -971,14 +971,16 @@ def intersect(self, x, y, z=None, local=False, forgive=False):
         When the point is on the edge of two cells, the cell with the lowest
         row or column is returned.
 
+        Supports both scalar and array inputs for vectorized operations.
+
         Parameters
         ----------
-        x : float
-            The x-coordinate of the requested point
-        y : float
-            The y-coordinate of the requested point
-        z : float
-            Optional z-coordinate of the requested point (will return layer,
+        x : float or array-like
+            The x-coordinate(s) of the requested point(s)
+        y : float or array-like
+            The y-coordinate(s) of the requested point(s)
+        z : float, array-like, or None
+            Optional z-coordinate(s) of the requested point(s) (will return layer,
             row, column) if supplied
         local: bool (optional)
             If True, x and y are in local coordinates (defaults to False)
@@ -988,59 +990,135 @@ def intersect(self, x, y, z=None, local=False, forgive=False):
 
         Returns
         -------
-        row : int
-            The row number
-        col : int
-            The column number
+        row : int or ndarray
+            The row number(s). Returns int for scalar input, ndarray for array input.
+        col : int or ndarray
+            The column number(s). Returns int for scalar input, ndarray for array input.
+        lay : int or ndarray (only if z is provided)
+            The layer number(s). Returns int for scalar input, ndarray for array input.
+
+        """
+        # Check if inputs are scalar
+        x_is_scalar = np.isscalar(x)
+        y_is_scalar = np.isscalar(y)
+        z_is_scalar = z is None or np.isscalar(z)
+        is_scalar_input = x_is_scalar and y_is_scalar and z_is_scalar
+
+        # Convert to arrays for uniform processing
+        x = np.atleast_1d(x)
+        y = np.atleast_1d(y)
+        if z is not None:
+            z = np.atleast_1d(z)
+
+        # Validate array shapes
+        if len(x) != len(y):
+            raise ValueError("x and y must have the same length")
+        if z is not None and len(z) != len(x):
+            raise ValueError("z must have the same length as x and y")
 
-        """
         # transform x and y to local coordinates
-        x, y = super().intersect(x, y, local, forgive)
+        if not local:
+            x, y = self.get_local_coords(x, y)
 
         # get the cell edges in local coordinates
         xe, ye = self.xyedges
 
-        xcomp = x > xe
-        if np.all(xcomp) or not np.any(xcomp):
-            if forgive:
-                col = np.nan
+        # Vectorized row/col calculation
+        n_points = len(x)
+        rows = np.full(n_points, np.nan if forgive else -1, dtype=float)
+        cols = np.full(n_points, np.nan if forgive else -1, dtype=float)
+
+        for i in range(n_points):
+            xi, yi = x[i], y[i]
+
+            # Find column
+            xcomp = xi > xe
+            if np.all(xcomp) or not np.any(xcomp):
+                if forgive:
+                    cols[i] = np.nan
+                else:
+                    raise ValueError(
+                        f"x, y point given is outside of the model area: ({xi}, {yi})"
+                    )
             else:
-                raise Exception("x, y point given is outside of the model area")
-        else:
-            col = np.asarray(xcomp).nonzero()[0][-1]
+                cols[i] = np.asarray(xcomp).nonzero()[0][-1]
 
-        ycomp = y < ye
-        if np.all(ycomp) or not np.any(ycomp):
-            if forgive:
-                row = np.nan
+            # Find row
+            ycomp = yi < ye
+            if np.all(ycomp) or not np.any(ycomp):
+                if forgive:
+                    rows[i] = np.nan
+                else:
+                    raise ValueError(
+                        f"x, y point given is outside of the model area: ({xi}, {yi})"
+                    )
             else:
-                raise Exception("x, y point given is outside of the model area")
-        else:
-            row = np.asarray(ycomp).nonzero()[0][-1]
-        if np.any(np.isnan([row, col])):
-            row = col = np.nan
-            if z is not None:
-                return None, row, col
+                rows[i] = np.asarray(ycomp).nonzero()[0][-1]
+
+        # If either row or col is NaN, set both to NaN
+        invalid_mask = np.isnan(rows) | np.isnan(cols)
+        rows[invalid_mask] = np.nan
+        cols[invalid_mask] = np.nan
+
+        # Convert to int where valid
+        valid_mask = ~invalid_mask
+        if np.any(valid_mask):
+            rows[valid_mask] = rows[valid_mask].astype(int)
+            cols[valid_mask] = cols[valid_mask].astype(int)
 
         if z is None:
-            return row, col
-
-        lay = np.nan
-        for layer in range(self.__nlay):
-            if (
-                self.top_botm[layer, row, col]
-                >= z
-                >= self.top_botm[layer + 1, row, col]
-            ):
-                lay = layer
-                break
-
-        if np.any(np.isnan([lay, row, col])):
-            lay = row = col = np.nan
-            if not forgive:
-                raise Exception("point given is outside the model area")
-
-        return lay, row, col
+            # Return results
+            if is_scalar_input:
+                row, col = rows[0], cols[0]
+                if not np.isnan(row) and not np.isnan(col):
+                    row, col = int(row), int(col)
+                return row, col
+            else:
+                return rows.astype(int) if np.all(valid_mask) else rows, cols.astype(
+                    int
+                ) if np.all(valid_mask) else cols
+
+        # Handle z-coordinate
+        lays = np.full(n_points, np.nan if forgive else -1, dtype=float)
+
+        for i in range(n_points):
+            if np.isnan(rows[i]) or np.isnan(cols[i]):
+                continue
+
+            row, col = int(rows[i]), int(cols[i])
+            zi = z[i]
+
+            for layer in range(self.__nlay):
+                if (
+                    self.top_botm[layer, row, col]
+                    >= zi
+                    >= self.top_botm[layer + 1, row, col]
+                ):
+                    lays[i] = layer
+                    break
+
+            if np.isnan(lays[i]) and not forgive:
+                raise ValueError(
+                    f"point given is outside the model area: ({x[i]}, {y[i]}, {zi})"
+                )
+
+        # Return results
+        if is_scalar_input:
+            lay, row, col = lays[0], rows[0], cols[0]
+            if not np.isnan(lay):
+                lay, row, col = int(lay), int(row), int(col)
+            else:
+                # When x,y are out of bounds: lay=None, row/col keep NaN
+                lay = None
+                # row and col already have their NaN values
+            return lay, row, col
+        else:
+            valid_3d = ~np.isnan(lays) & ~np.isnan(rows) & ~np.isnan(cols)
+            return (
+                lays.astype(int) if np.all(valid_3d) else lays,
+                rows.astype(int) if np.all(valid_3d) else rows,
+                cols.astype(int) if np.all(valid_3d) else cols,
+            )
 
     def _cell_vert_list(self, i, j):
         """Get vertices for a single cell or sequence of i, j locations."""