diff --git a/autotest/conftest.py b/autotest/conftest.py index 036ce9e67..53650cc00 100644 --- a/autotest/conftest.py +++ b/autotest/conftest.py @@ -7,6 +7,8 @@ import pytest from modflow_devtools.misc import is_in_ci +from autotest.test_grid_cases import GridCases + # import modflow-devtools fixtures pytest_plugins = ["modflow_devtools.fixtures", "modflow_devtools.snapshots"] @@ -79,6 +81,51 @@ def patch_macos_ci_matplotlib(): matplotlib.use("agg") +# grid fixtures (from GridCases) + + +@pytest.fixture +def minimal_vertex_grid(): + """Minimal 2-cell vertex grid.""" + return GridCases.vertex_minimal() + + +@pytest.fixture +def minimal_unstructured_grid(): + """Minimal 2-cell unstructured grid.""" + return GridCases.unstructured_minimal() + + +@pytest.fixture +def triangular_vertex_grid(): + """Triangular vertex grid using Delaunay triangulation.""" + return GridCases.vertex_triangular() + + +@pytest.fixture +def triangular_unstructured_grid(): + """Triangular unstructured grid using Delaunay triangulation.""" + return GridCases.unstructured_triangular() + + +@pytest.fixture +def simple_vertex_grid(): + """10x10 rectangular vertex grid (100 cells, 10x10 each).""" + return GridCases.vertex_rectangular(nrow=10, ncol=10, cell_size=10.0) + + +@pytest.fixture +def rotated_vertex_grid(): + """Rotated 5x5 rectangular vertex grid.""" + return GridCases.vertex_rectangular(nrow=5, ncol=5, cell_size=20.0, angrot=45.0) + + +@pytest.fixture +def layered_unstructured_grid(): + """3-layer unstructured grid for 3D testing.""" + return GridCases.unstructured_layered() + + # pytest configuration hooks diff --git a/autotest/test_grid.py b/autotest/test_grid.py index 6f1211e60..e9c0fd84f 100644 --- a/autotest/test_grid.py +++ b/autotest/test_grid.py @@ -26,6 +26,7 @@ gridlist_to_disv_gridprops, to_cvfd, ) +from flopy.utils.geospatial_index import GeospatialIndex from flopy.utils.triangle import Triangle from flopy.utils.voronoi import VoronoiGrid @@ -1603,6 +1604,53 @@ def test_geo_dataframe(structured_grid, vertex_grid, unstructured_grid): raise AssertionError(f"Cell vertices incorrect for node={node}") +def test_structured_boundary_tiebreaker(simple_structured_grid): + """Test StructuredGrid uses lowest row/col for boundary points.""" + grid = simple_structured_grid + + # Boundary at x=10 -> col 0 + _, col = grid.intersect(10.0, 95.0) + assert col == 0 + + # Boundary at y=90 -> row 0 + row, _ = grid.intersect(5.0, 90.0) + assert row == 0 + + +def test_return_type_structured(simple_structured_grid): + """Test StructuredGrid return types are consistent.""" + grid = simple_structured_grid + + # Scalar inside -> int + row, col = grid.intersect(5.0, 95.0) + assert isinstance(row, (int, np.integer)) + assert isinstance(col, (int, np.integer)) + + # Scalar outside -> nan + row, col = grid.intersect(150.0, 50.0, forgive=True) + assert np.isnan(row) and np.isnan(col) + + # Scalar with z inside -> int + lay, row, col = grid.intersect(5.0, 95.0, z=5.0) + assert all(isinstance(v, (int, np.integer)) for v in [lay, row, col]) + + # Scalar with z outside -> nan + lay, row, col = grid.intersect(150.0, 50.0, z=5.0, forgive=True) + assert all(np.isnan(v) for v in [lay, row, col]) + + # Array all inside -> any integer dtype + rows, cols = grid.intersect(np.array([5.0, 55.0]), np.array([95.0, 95.0])) + assert np.issubdtype(rows.dtype, np.integer) + assert np.issubdtype(cols.dtype, np.integer) + + # Array with outside -> must be float64 (due to NaNs) + rows, cols = grid.intersect( + np.array([5.0, 150.0]), np.array([95.0, 50.0]), forgive=True + ) + assert rows.dtype == np.float64 + assert cols.dtype == np.float64 + + def test_unstructured_iverts_cleanup(): grid = GridCases.structured_small() @@ -1657,18 +1705,236 @@ def test_unstructured_iverts_cleanup(): 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. +# ============================================================================ +# GeospatialIndex-specific Tests +# ============================================================================ + + +@pytest.mark.parametrize( + "grid_fixture,grid_type,expected_cells,expected_points_per_cell", + [ + ("minimal_vertex_grid", "vertex", 2, 5), + ("minimal_unstructured_grid", "unstructured", 2, 5), + ("simple_vertex_grid", "vertex", 100, 5), + ], +) +def test_index_build( + grid_fixture, grid_type, expected_cells, expected_points_per_cell, request +): + """Test building GeospatialIndex for different grid types.""" + grid = request.getfixturevalue(grid_fixture) + index = GeospatialIndex(grid) + + assert index.grid.grid_type == grid_type + assert index.grid is grid + assert index.tree is not None + assert index.points.shape[0] == expected_cells * expected_points_per_cell + assert len(np.unique(index.point_to_cell)) == expected_cells + + +def test_index_repr(simple_vertex_grid): + """Test string representation of GeospatialIndex.""" + index = GeospatialIndex(simple_vertex_grid) + repr_str = repr(index) + + assert "GeospatialIndex" in repr_str + assert "vertex grid" in repr_str + assert "100 cells" in repr_str + assert "500 indexed points" in repr_str + + +@pytest.mark.parametrize("k", [1, 5, 10, 20]) +def test_index_k_parameter(simple_vertex_grid, k): + """Test that different k values still find correct cell.""" + index = GeospatialIndex(simple_vertex_grid) + + assert index.query_point(5.0, 95.0, k=k) == 0 + + cellids = index.query_points( + np.array([5.0, 55.0, 95.0]), + np.array([95.0, 95.0, 5.0]), + k=k, + ) + np.testing.assert_array_equal(cellids, [0, 5, 99]) - 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). + +@pytest.fixture +def sliver_grid_geometry(): + """Shared geometry for sliver cell tests. + + Creates a 1000:1 aspect ratio sliver cell (100 units wide, 0.1 units tall) + completely surrounded by larger cells on all four sides. This is the + worst-case scenario for centroid-based spatial indexing because query + points near the sliver edges are much closer to neighboring cell centroids. + + Returns dict with vertices, iverts, xcenters, ycenters for 5 cells: + Cell 0: SLIVER (x: 0-100, y: 0-0.1) + Cell 1: TOP (above sliver) + Cell 2: BOTTOM (below sliver) + Cell 3: LEFT (left of sliver) + Cell 4: RIGHT (right of sliver) """ - grid = simple_structured_grid + sliver_height = 0.1 + sliver_yc = sliver_height / 2 + top_yc = (sliver_height + 10.0) / 2 + + # Base x,y coordinates for vertices (no vertex IDs yet) + base_vertices = [ + # Cell 0: SLIVER + (0.0, 0.0), + (100.0, 0.0), + (100.0, sliver_height), + (0.0, sliver_height), + # Cell 1: TOP + (0.0, sliver_height), + (100.0, sliver_height), + (100.0, 10.0), + (0.0, 10.0), + # Cell 2: BOTTOM + (0.0, -10.0), + (100.0, -10.0), + (100.0, 0.0), + (0.0, 0.0), + # Cell 3: LEFT + (-10.0, -10.0), + (0.0, -10.0), + (0.0, 10.0), + (-10.0, 10.0), + # Cell 4: RIGHT + (100.0, -10.0), + (110.0, -10.0), + (110.0, 10.0), + (100.0, 10.0), + ] + + base_iverts = [ + [0, 1, 2, 3], + [4, 5, 6, 7], + [8, 9, 10, 11], + [12, 13, 14, 15], + [16, 17, 18, 19], + ] + + xcenters = [50.0, 50.0, 50.0, -5.0, 105.0] + ycenters = [sliver_yc, top_yc, -5.0, 0.0, 0.0] + + return { + "base_vertices": base_vertices, + "base_iverts": base_iverts, + "xcenters": xcenters, + "ycenters": ycenters, + "sliver_height": sliver_height, + } + + +def test_index_thin_sliver_cell(sliver_grid_geometry): + """Test GeospatialIndex finds points in very thin sliver cells (2D). + + Query point at (95, 0.05) inside the sliver: + - Distance to sliver centroid (50, 0.05): 45 units + - Distance to RIGHT cell centroid (105, 0): ~10 units <- MUCH CLOSER! + + A centroid-only KD-tree would return the wrong cell. + """ + geom = sliver_grid_geometry + + # Build VertexGrid vertices with IDs + vertices = [[i, x, y] for i, (x, y) in enumerate(geom["base_vertices"])] + + # Build cell2d from geometry + cell2d = [] + for cid, iverts in enumerate(geom["base_iverts"]): + xc = geom["xcenters"][cid] + yc = geom["ycenters"][cid] + cell2d.append([cid, xc, yc, len(iverts)] + iverts) + + ncells = len(cell2d) + grid = VertexGrid( + vertices=vertices, + cell2d=cell2d, + top=np.ones(ncells) * 10.0, + botm=np.zeros(ncells), + ) + + index = GeospatialIndex(grid) + + # Test sliver cell detection at edges (where neighboring centroids are closer) + assert index.query_point(95.0, 0.05, k=10) == 0, "Right edge of sliver" + assert index.query_point(5.0, 0.05, k=10) == 0, "Left edge of sliver" + assert index.query_point(50.0, 0.05, k=10) == 0, "Center of sliver" + + # Sanity check surrounding cells + assert index.query_point(50.0, 5.0, k=10) == 1, "Top cell" + assert index.query_point(50.0, -5.0, k=10) == 2, "Bottom cell" + assert index.query_point(-5.0, 0.0, k=10) == 3, "Left cell" + assert index.query_point(105.0, 0.0, k=10) == 4, "Right cell" + + +def test_index_thin_sliver_cell_3d(sliver_grid_geometry): + """Test GeospatialIndex finds thin sliver cells in true 3D grids. + + Uses UnstructuredGrid with grid_varies_by_layer=True to test the 3D + KD-tree and 3D bounding box lookup with the same x-y sliver geometry. + """ + geom = sliver_grid_geometry + nlay = 3 + ncells_per_layer = len(geom["base_iverts"]) + nverts_per_layer = len(geom["base_vertices"]) + + # Replicate geometry across layers for grid_varies_by_layer=True + vertices = [] + iverts = [] + xcenters = [] + ycenters = [] + + for lay in range(nlay): + vert_offset = lay * nverts_per_layer + for i, (x, y) in enumerate(geom["base_vertices"]): + vertices.append([vert_offset + i, x, y]) + for cell_iverts in geom["base_iverts"]: + iverts.append([v + vert_offset for v in cell_iverts]) + xcenters.extend(geom["xcenters"]) + ycenters.extend(geom["ycenters"]) + + ncpl = np.array([ncells_per_layer] * nlay) + nnodes = np.sum(ncpl) + + # Set up layer elevations + top = np.zeros(nnodes) + botm = np.zeros(nnodes) + layer_tops = [100.0, 66.67, 33.33] + layer_bots = [66.67, 33.33, 0.0] + + for lay in range(nlay): + start = lay * ncells_per_layer + end = start + ncells_per_layer + top[start:end] = layer_tops[lay] + botm[start:end] = layer_bots[lay] + + grid = UnstructuredGrid( + vertices=vertices, + iverts=iverts, + xcenters=xcenters, + ycenters=ycenters, + ncpl=ncpl, + top=top, + botm=botm, + ) + + assert grid.grid_varies_by_layer, "Grid should have grid_varies_by_layer=True" + + index = GeospatialIndex(grid) + assert index.is_3d, "Index should be in 3D mode" + + # Test sliver cell in each layer (cell 0, 5, 10 are slivers) + # Point (95, 0.05) is far from sliver centroid but inside the cell + assert index.intersect(95.0, 0.05, z=83.33) == 0, "Sliver layer 0" + assert index.intersect(95.0, 0.05, z=50.0) == 5, "Sliver layer 1" + assert index.intersect(95.0, 0.05, z=16.67) == 10, "Sliver layer 2" - # Test out-of-bounds x,y with z coordinate - lay, row, col = grid.intersect(-50.0, -50.0, z=5.0, forgive=True) + # Test left edge of sliver + assert index.intersect(5.0, 0.05, z=83.33) == 0, "Left edge layer 0" - # 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}" + # Test surrounding cells + assert index.intersect(105.0, 0.0, z=83.33) == 4, "Right cell layer 0" + assert index.intersect(50.0, 5.0, z=50.0) == 6, "Top cell layer 1" diff --git a/autotest/test_grid_cases.py b/autotest/test_grid_cases.py index 9415a237b..12432002e 100644 --- a/autotest/test_grid_cases.py +++ b/autotest/test_grid_cases.py @@ -2,6 +2,7 @@ from tempfile import TemporaryDirectory import numpy as np +from scipy.spatial import Delaunay from flopy.discretization import StructuredGrid, UnstructuredGrid, VertexGrid from flopy.utils.triangle import Triangle @@ -387,3 +388,167 @@ def voronoi_many_polygons(): grid = VertexGrid(**gridprops, nlay=1) return ncpl, vor, gridprops, grid + + # ========================================================================= + # Minimal grids for GeospatialIndex testing + # ========================================================================= + + @staticmethod + def _minimal_geometry(): + """Create shared 2-cell geometry used by multiple grid methods.""" + vertices = [ + [0, 0.0, 1.0], + [1, 1.0, 1.0], + [2, 2.0, 1.0], + [3, 0.0, 0.0], + [4, 1.0, 0.0], + [5, 2.0, 0.0], + ] + iverts = [[0, 1, 4, 3], [1, 2, 5, 4]] + xcenters = [0.5, 1.5] + ycenters = [0.5, 0.5] + return vertices, iverts, xcenters, ycenters + + @staticmethod + def _triangular_geometry(seed=42, n_points=30): + """Create shared Delaunay triangulation geometry.""" + np.random.seed(seed) + 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] + + return vertices, iverts, xcenters, ycenters + + @staticmethod + def vertex_minimal(): + """Create a minimal 2-cell vertex grid.""" + vertices, iverts, xcenters, ycenters = GridCases._minimal_geometry() + cell2d = [ + [0, xcenters[0], ycenters[0], 4] + iverts[0], + [1, xcenters[1], ycenters[1], 4] + iverts[1], + ] + return VertexGrid(vertices=vertices, cell2d=cell2d, nlay=1) + + @staticmethod + def unstructured_minimal(): + """Create a minimal 2-cell unstructured grid.""" + vertices, iverts, xcenters, ycenters = GridCases._minimal_geometry() + return UnstructuredGrid( + vertices=vertices, iverts=iverts, xcenters=xcenters, ycenters=ycenters + ) + + @staticmethod + def vertex_triangular(seed=42, n_points=30): + """Create a triangular vertex grid using Delaunay triangulation.""" + vertices, iverts, xcenters, ycenters = GridCases._triangular_geometry( + seed, n_points + ) + cell2d = [ + [i, xcenters[i], ycenters[i], 3] + iverts[i] for i in range(len(iverts)) + ] + ncells = len(cell2d) + return VertexGrid( + vertices=vertices, + cell2d=cell2d, + top=np.ones(ncells) * 10.0, + botm=np.zeros(ncells), + ) + + @staticmethod + def unstructured_triangular(seed=42, n_points=30): + """Create a triangular unstructured grid using Delaunay triangulation.""" + vertices, iverts, xcenters, ycenters = GridCases._triangular_geometry( + seed, n_points + ) + ncells = len(iverts) + return UnstructuredGrid( + vertices=vertices, + iverts=iverts, + xcenters=xcenters, + ycenters=ycenters, + top=np.ones(ncells) * 10.0, + botm=np.zeros(ncells), + ) + + @staticmethod + def vertex_rectangular(nrow=10, ncol=10, cell_size=10.0, angrot=0.0): + """Create a rectangular vertex grid. + + Parameters + ---------- + nrow : int + Number of rows. + ncol : int + Number of columns. + cell_size : float + Size of each cell (square cells). + angrot : float + Grid rotation angle in degrees. + """ + vertices = [] + vid = 0 + for i in range(nrow + 1): + for j in range(ncol + 1): + x = j * cell_size + y = (nrow - i) * cell_size + vertices.append([vid, x, y]) + vid += 1 + + cell2d = [] + cellid = 0 + for i in range(nrow): + for j in range(ncol): + xc = (j + 0.5) * cell_size + yc = (nrow - i - 0.5) * cell_size + v0 = i * (ncol + 1) + j + v1 = v0 + 1 + v2 = v1 + (ncol + 1) + v3 = v0 + (ncol + 1) + cell2d.append([cellid, xc, yc, 4, v0, v1, v2, v3]) + cellid += 1 + + top = np.ones(nrow * ncol) * 10.0 + botm = np.zeros(nrow * ncol) + + return VertexGrid( + vertices=vertices, cell2d=cell2d, top=top, botm=botm, nlay=1, angrot=angrot + ) + + @staticmethod + def unstructured_layered(): + """Create a 3-layer unstructured grid for 3D testing. + + Grid has 2 cells per layer, 3 layers total = 6 cells. + Cell IDs: 0-1 (layer 0), 2-3 (layer 1), 4-5 (layer 2). + Z elevations: layer 0 (10-7), layer 1 (7-4), layer 2 (4-1). + + Uses 2D indexing with layer search (grid_varies_by_layer=False). + """ + vertices, base_iverts, base_xc, base_yc = GridCases._minimal_geometry() + + nlay = 3 + ncpl = 2 + + # Only provide iverts for first layer (grid_varies_by_layer=False) + iverts = base_iverts + xcenters = list(base_xc) + ycenters = list(base_yc) + + top = np.array([10.0, 10.0, 7.0, 7.0, 4.0, 4.0]) + botm = np.array([7.0, 7.0, 4.0, 4.0, 1.0, 1.0]) + + return UnstructuredGrid( + vertices=vertices, + iverts=iverts, + xcenters=xcenters, + ycenters=ycenters, + top=top, + botm=botm, + ncpl=np.array([ncpl] * nlay), + ) diff --git a/flopy/discretization/grid.py b/flopy/discretization/grid.py index 1d59801b0..a4e560167 100644 --- a/flopy/discretization/grid.py +++ b/flopy/discretization/grid.py @@ -993,11 +993,58 @@ def get_local_coords(self, x, y): return x, y - def intersect(self, x, y, local=False, forgive=False): - if not local: - return self.get_local_coords(x, y) - else: - return x, y + def intersect(self, x, y, z=None, local=False, forgive=False): + """ + Get the cell(s) containing the given point(s). + + Uses GeospatialIndex for efficient spatial queries, with optimal + algorithms for each grid type: + - StructuredGrid: searchsorted (O(log n)) + - VertexGrid/UnstructuredGrid: KD-tree + ConvexHull + + When the point is on the edge of two cells, the cell with the lowest + index is returned. + + Supports both scalar and array inputs for vectorized operations. + + Parameters + ---------- + x : float or array-like + The x-coordinate(s) of the query point(s) + y : float or array-like + The y-coordinate(s) of the query point(s) + z : float, array-like, or None + Optional z-coordinate(s). If provided, returns layer information. + local : bool, optional + If True, x and y are in local coordinates (default False) + forgive : bool, optional + If True, return NaN for points outside the grid instead of + raising an error (default False) + + Returns + ------- + For StructuredGrid: + row, col : int or ndarray + Row and column indices. If z is provided, returns (lay, row, col). + For VertexGrid: + cellid : int or ndarray + Cell index (icell2d). If z is provided, returns (lay, cellid). + For UnstructuredGrid: + cellid : int or ndarray + Cell index. + + Raises + ------ + ValueError + If point is outside grid and forgive=False + """ + from ..utils.geospatial_index import GeospatialIndex + + # Lazily create geospatial index + if not hasattr(self, "_geospatial_index") or self._geospatial_index is None: + self._geospatial_index = GeospatialIndex(self) + + return self._geospatial_index.intersect(x, y, z=z, local=local, forgive=forgive) def set_coord_info( self, diff --git a/flopy/discretization/structuredgrid.py b/flopy/discretization/structuredgrid.py index 0d2df58fd..43badef68 100644 --- a/flopy/discretization/structuredgrid.py +++ b/flopy/discretization/structuredgrid.py @@ -964,162 +964,6 @@ def _fast_neighbors(self, **kwargs): return self._neighbors - def intersect(self, x, y, z=None, local=False, forgive=False): - """ - Get the row and column of a point with coordinates x and y - - 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 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) - forgive: bool (optional) - Forgive x,y arguments that fall outside the model grid and - return NaNs instead (defaults to False - will throw exception) - - Returns - ------- - 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 - if not local: - x, y = self.get_local_coords(x, y) - - # get the cell edges in local coordinates - xe, ye = self.xyedges - - # 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: - cols[i] = np.asarray(xcomp).nonzero()[0][-1] - - # 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: - 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 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.""" self._copy_cache = False diff --git a/flopy/discretization/unstructuredgrid.py b/flopy/discretization/unstructuredgrid.py index f34764dd6..c69f15786 100644 --- a/flopy/discretization/unstructuredgrid.py +++ b/flopy/discretization/unstructuredgrid.py @@ -25,13 +25,17 @@ class UnstructuredGrid(Grid): size nodes, if the grid_varies_by_nodes argument is true, or it must be of size ncpl[0] if the same 2d spatial grid is used for each layer. xcenters : list or ndarray - list of x center coordinates for all cells in the grid if the grid - varies by layer or for all cells in a layer if the same grid is used - for all layers + x-coordinates of cell centers for all cells in the grid if the grid + varies by layer, or for all cells in a layer if the same grid is used + for all layers. These are typically geometric centroids but may be + any representative point. For convex cells (triangles, rectangles), + the arithmetic mean of vertices is equivalent to the true centroid. + For concave cells, the provided center should ideally lie inside the + cell boundary for optimal spatial indexing performance. ycenters : list or ndarray - list of y center coordinates for all cells in the grid if the grid - varies by layer or for all cells in a layer if the same grid is used - for all layers + y-coordinates of cell centers for all cells in the grid if the grid + varies by layer, or for all cells in a layer if the same grid is used + for all layers. See ``xcenters`` for details on center point placement. top : list or ndarray top elevations for all cells in the grid. botm : list or ndarray @@ -718,129 +722,6 @@ def clean_iverts(self, inplace=False): ja=self._ja, ) - def intersect(self, x, y, z=None, local=False, forgive=False): - """ - Get the CELL2D number of a point with coordinates x and y - - When the point is on the edge of two cells, the cell with the lowest - CELL2D number is returned. - - Supports both scalar and array inputs for vectorized operations. - - Parameters - ---------- - 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) - local: bool (optional) - If True, x and y are in local coordinates (defaults to False) - forgive: bool (optional) - Forgive x,y arguments that fall outside the model grid and - return NaNs instead (defaults to False - will throw exception) - - Returns - ------- - icell2d : int or ndarray - The CELL2D 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") - - if local: - # transform x and y to real-world coordinates - x, y = super().get_coords(x, y) - - xv, yv, zv = self.xyzvertices - - if self.grid_varies_by_layer: - ncpl = self.nnodes - else: - ncpl = self.ncpl[0] - - # Initialize result array - n_points = len(x) - results = np.full(n_points, np.nan if forgive else -1, dtype=float) - - # Process each point - for i in range(n_points): - xi, yi = x[i], y[i] - zi = z[i] if z is not None else None - found = False - - for icell2d in range(ncpl): - xa = np.array(xv[icell2d]) - ya = np.array(yv[icell2d]) - # x and y at least have to be within the bounding box of the cell - if ( - np.any(xi <= xa) - and np.any(xi >= xa) - and np.any(yi <= ya) - and np.any(yi >= ya) - ): - if is_clockwise(xa, ya): - radius = -1e-9 - else: - radius = 1e-9 - path = Path(np.stack((xa, ya)).transpose()) - # use a small radius, so that the edge of the cell is included - if path.contains_point((xi, yi), radius=radius): - if zi is None: - results[i] = icell2d - found = True - break - - # Search through layers for z-coordinate - cell_idx_3d = icell2d - for lay in range(self.nlay): - if zv[0, cell_idx_3d] >= zi >= zv[1, cell_idx_3d]: - results[i] = cell_idx_3d - found = True - break - # Move to next layer - if lay < self.nlay - 1 and not self.grid_varies_by_layer: - cell_idx_3d += self.ncpl[lay] - if found: - break - - if not found and not forgive: - raise ValueError(f"point ({xi}, {yi}) is outside of the model area") - - # Return scalar if input was scalar, otherwise return array - if is_scalar_input: - result = results[0] - return int(result) if not np.isnan(result) else np.nan - else: - # Convert to int array where not NaN - if not forgive: - return results.astype(int) - else: - # Keep as float to preserve NaN values - valid_mask = ~np.isnan(results) - if np.all(valid_mask): - return results.astype(int) - else: - return results - @property def top_botm(self): new_top = np.expand_dims(self._top, 0) diff --git a/flopy/discretization/vertexgrid.py b/flopy/discretization/vertexgrid.py index dc5be8710..acef1b5f1 100644 --- a/flopy/discretization/vertexgrid.py +++ b/flopy/discretization/vertexgrid.py @@ -15,9 +15,16 @@ class for a vertex model grid Parameters ---------- vertices - list of vertices that make up the grid + list of vertices that make up the grid. Each vertex is + [iv, xv, yv] where iv is the vertex number. cell2d - list of cells and their vertices + list of cells and their vertices. Each cell is + [icell2d, xc, yc, ncvert, icvert1, icvert2, ...] where xc, yc are + cell center coordinates. These are typically geometric centroids but + may be any representative point inside the cell. For convex cells + (triangles, rectangles), the arithmetic mean of vertices equals the + true centroid. For concave cells, the center should ideally lie + inside the cell boundary for optimal spatial indexing performance. top : list or ndarray top elevations for all cells in the grid. botm : list or ndarray @@ -344,136 +351,6 @@ def convert_grid(self, factor): else: raise AssertionError("Grid is not complete and cannot be converted") - def intersect(self, x, y, z=None, local=False, forgive=False): - """ - Get the CELL2D number of a point with coordinates x and y - - When the point is on the edge of two cells, the cell with the lowest - CELL2D number is returned. - - Supports both scalar and array inputs for vectorized operations. - - Parameters - ---------- - 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 - (lay, icell2d) - local: bool (optional) - If True, x and y are in local coordinates (defaults to False) - forgive: bool (optional) - Forgive x,y arguments that fall outside the model grid and - return NaNs instead (defaults to False - will throw exception) - - Returns - ------- - icell2d : int or ndarray - The CELL2D 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") - - if local: - # transform x and y to real-world coordinates - x, y = super().get_coords(x, y) - - xv, yv, zv = self.xyzvertices - - # Initialize result arrays - n_points = len(x) - results = np.full(n_points, np.nan if forgive else -1, dtype=float) - if z is not None: - lays = np.full(n_points, np.nan if forgive else -1, dtype=float) - - # Process each point - for i in range(n_points): - xi, yi = x[i], y[i] - found = False - - # Check each cell - for icell2d in range(self.ncpl): - xa = np.array(xv[icell2d]) - ya = np.array(yv[icell2d]) - # x and y at least have to be within the bounding box of the cell - if ( - np.any(xi <= xa) - and np.any(xi >= xa) - and np.any(yi <= ya) - and np.any(yi >= ya) - ): - path = Path(np.stack((xa, ya)).transpose()) - # use a small radius, so that the edge of the cell is included - if is_clockwise(xa, ya): - radius = -1e-9 - else: - radius = 1e-9 - if path.contains_point((xi, yi), radius=radius): - results[i] = icell2d - found = True - - if z is not None: - zi = z[i] - for lay in range(self.nlay): - if ( - self.top_botm[lay, icell2d] - >= zi - >= self.top_botm[lay + 1, icell2d] - ): - lays[i] = lay - break - - break - - if not found and not forgive: - raise ValueError( - f"point given is outside of the model area: ({xi}, {yi})" - ) - - # Return results - if z is None: - if is_scalar_input: - result = results[0] - return int(result) if not np.isnan(result) else np.nan - else: - valid_mask = ~np.isnan(results) - return results.astype(int) if np.all(valid_mask) else results - else: - if is_scalar_input: - lay, icell2d = lays[0], results[0] - if not np.isnan(lay) and not np.isnan(icell2d): - return int(lay), int(icell2d) - else: - return np.nan, np.nan - else: - valid_mask = ~np.isnan(lays) & ~np.isnan(results) - return ( - lays.astype(int) if np.all(valid_mask) else lays, - results.astype(int) if np.all(valid_mask) else results, - ) - def get_cell_vertices(self, cellid): """ Method to get a set of cell vertices for a single cell diff --git a/flopy/utils/geospatial_index.py b/flopy/utils/geospatial_index.py new file mode 100644 index 000000000..7e09ffc9c --- /dev/null +++ b/flopy/utils/geospatial_index.py @@ -0,0 +1,964 @@ +""" +Geospatial indexing for FloPy grids. + +Provides efficient spatial queries for all grid types: +- StructuredGrid: Uses searchsorted for O(log n) row/column finding +- VertexGrid/UnstructuredGrid: Uses KD-tree with cell centers + vertices, + plus pre-computed ConvexHull equations for fast point-in-polygon testing. + +This module provides a unified spatial query interface for all grid types, +with each type using the optimal algorithm for its geometry. +""" + +import numpy as np +from scipy.spatial import ConvexHull, cKDTree + + +class GeospatialIndex: + """ + Geospatial index for efficient spatial queries on any FloPy grid. + + Provides a unified interface for point intersection queries across all + grid types, using the optimal algorithm for each: + + - **StructuredGrid**: Uses numpy searchsorted for O(log n) row/column + finding. No index structure is built; queries operate directly on + the grid's edge arrays. + + - **VertexGrid/UnstructuredGrid**: Uses KD-tree indexing with cell + centers + vertices to find candidate cells, then pre-computed + ConvexHull hyperplane equations for fast vectorized point-in-polygon + testing. + + The cell center + vertices approach for unstructured grids ensures + edge cases are handled: + - Points near cell boundaries + - Points in thin/sliver cells where the cell center may be far from the query + + Note + ---- + For vertex/unstructured grids, this index uses the grid's ``xcellcenters`` + and ``ycellcenters`` properties, which represent user-provided or computed + cell center coordinates. These are not necessarily true geometric centroids + (center of mass). The index handles this by also indexing all cell vertices, + ensuring robust spatial queries regardless of cell center placement. + + Parameters + ---------- + grid : StructuredGrid, VertexGrid, or UnstructuredGrid + A FloPy grid object + epsilon : float, optional + Tolerance for point-in-cell tests. Used for both bounding box + expansion and ConvexHull hyperplane distance tests. Ensures boundary + points are included in adjacent cells. + + Attributes + ---------- + grid : Grid + The grid object this index was built for + epsilon : float + Tolerance for geometric tests + is_3d : bool + True if index uses 3D coordinates (x,y,z), False for 2D (x,y only). + Automatically True when grid has grid_varies_by_layer=True. + tree : scipy.spatial.cKDTree or None + KD-tree of cell centers + vertices for fast spatial queries + (vertex/unstructured grids only). None for structured grids. + point_to_cell : ndarray or None + Mapping from KD-tree point index to cell index + (vertex/unstructured grids only). None for structured grids. + hull_equations : list or None + Pre-computed ConvexHull equations for each cell (2D vertex/unstructured + grids only). None for structured grids. + bounding_boxes : ndarray or None + Pre-computed bounding boxes for each cell (vertex/unstructured + grids only). None for structured grids. + + Examples + -------- + >>> from flopy.discretization import StructuredGrid, VertexGrid + >>> from flopy.utils.geospatial_index import GeospatialIndex + >>> + >>> # Works with any grid type + >>> index = GeospatialIndex(grid) + >>> + >>> # Single point query + >>> cellid = index.query_point(x=5.5, y=5.5) + >>> + >>> # Multiple points (vectorized) + >>> cellids = index.query_points(x=[1.5, 5.5, 9.5], y=[1.5, 5.5, 9.5]) + >>> + >>> # For structured grids, get (row, col) tuples + >>> row, col = index.intersect(x=5.5, y=5.5) + """ + + def __init__(self, grid, epsilon=1e-3): + """ + Build geospatial index for any FloPy grid type. + + Parameters + ---------- + grid : StructuredGrid, VertexGrid, or UnstructuredGrid + A FloPy grid object + epsilon : float, optional + Tolerance for point-in-cell tests. + Used for bounding box expansion and ConvexHull tests. + """ + self.grid = grid + self.epsilon = epsilon + + # Initialize attributes that may not be set for all grid types + self.tree = None + self.points = None + self.point_to_cell = None + self.hull_equations = None + self.bounding_boxes = None + self.bounding_boxes_3d = None + + self._build_index() + + @property + def is_3d(self): + """True if grid geometry varies by layer (requires 3D indexing).""" + return getattr(self.grid, "grid_varies_by_layer", False) + + def _build_index(self): + """ + Build spatial index appropriate for the grid type. + + For structured grids: No index structure needed; uses searchsorted + directly on edge arrays. + + For vertex/unstructured grids: + - 2D: KD-tree with cell centers + vertices, ConvexHull equations + - 3D (grid_varies_by_layer=True): 3D KD-tree with bounding boxes + """ + # Structured grids don't need an index structure + if self.grid.grid_type == "structured": + self.ncells = self.grid.nnodes + # Cache edge arrays for faster queries + self._xe, self._ye = self.grid.xyedges + self._ye_flipped = self._ye[::-1] + return + + # Vertex/unstructured grids use KD-tree indexing + points = [] + point_to_cell = [] + + if self.is_3d: + # 3D indexing: index all nnodes with x,y,z coordinates + self.ncells = self.grid.nnodes + self._build_3d_index(points, point_to_cell) + else: + # 2D indexing: index only 2D cells with x,y coordinates + if hasattr(self.grid, "ncpl"): + ncpl = self.grid.ncpl + if ncpl is None: + # ncpl not set, fall back to xcellcenters + self.ncells = len(self.grid.xcellcenters) + elif np.isscalar(ncpl): + # VertexGrid: ncpl is scalar + self.ncells = ncpl + else: + # UnstructuredGrid: ncpl is array + # Use first layer's cell count for 2D spatial indexing + if len(ncpl) > 0: + self.ncells = ncpl[0] + else: + self.ncells = len(self.grid.xcellcenters) + else: + self.ncells = len(self.grid.xcellcenters) + self._build_vertex_index(points, point_to_cell) + + # Build KD-tree from cell centers + vertices + self.points = np.array(points) + self.point_to_cell = np.array(point_to_cell, dtype=int) + self.tree = cKDTree(self.points) + + # Pre-compute ConvexHull equations and bounding boxes (2D only) + if not self.is_3d: + self._precompute_hulls() + else: + # For 3D, store z-bounds for each cell + self._precompute_3d_bounds() + + def _build_vertex_index(self, points, point_to_cell): + """Build index for vertex/unstructured grid - cell centers + vertices.""" + # Disable copy cache to avoid expensive deepcopy during index build + original_copy_cache = getattr(self.grid, "_copy_cache", True) + self.grid._copy_cache = False + + xc = self.grid.xcellcenters + yc = self.grid.ycellcenters + xv, yv, _ = self.grid.xyzvertices + + for cellid in range(self.ncells): + # Add cell center + points.append([xc[cellid], yc[cellid]]) + point_to_cell.append(cellid) + + # Add all cell vertices + cell_xv = xv[cellid] + cell_yv = yv[cellid] + for vi in range(len(cell_xv)): + points.append([cell_xv[vi], cell_yv[vi]]) + point_to_cell.append(cellid) + + # Restore copy cache setting + self.grid._copy_cache = original_copy_cache + + def _build_3d_index(self, points, point_to_cell): + """Build 3D index for grid_varies_by_layer grids. + + Includes cell centers + vertices with z-coordinates. + """ + # Disable copy cache to avoid expensive deepcopy during index build + original_copy_cache = getattr(self.grid, "_copy_cache", True) + self.grid._copy_cache = False + + xc = np.asarray(self.grid.xcellcenters) + yc = np.asarray(self.grid.ycellcenters) + xv, yv, zv = self.grid.xyzvertices + + # Get z-coordinates for cell centers (use mid-point of top/bottom) + zc = (zv[0] + zv[1]) / 2.0 + + # For multi-layer unstructured grids, xc/yc may be per-layer (ncpl values) + # while we iterate over all nnodes. Build mapping from node to layer-local cell. + ncpl = getattr(self.grid, "ncpl", None) + if ncpl is not None and not np.isscalar(ncpl): + ncpl = np.atleast_1d(ncpl) + xc_is_per_layer = len(xc) < self.ncells + else: + xc_is_per_layer = False + + # Build xc_idx mapping vectorized + cellids = np.arange(self.ncells) + if xc_is_per_layer: + cumulative_ncpl = np.concatenate([[0], np.cumsum(ncpl[:-1])]) + layers = np.searchsorted(np.cumsum(ncpl), cellids, side="right") + xc_idx = cellids - cumulative_ncpl[layers] + else: + xc_idx = cellids + + # Add cell centers + centers = np.column_stack([xc[xc_idx], yc[xc_idx], zc]) + center_cells = cellids + + # Add vertices - must loop due to variable vertex count per cell + vertex_points = [] + vertex_cells = [] + zc_mid = (zv[0] + zv[1]) / 2.0 + for cellid in range(self.ncells): + cell_xv = np.atleast_1d(xv[cellid]) + cell_yv = np.atleast_1d(yv[cellid]) + cell_z = zc_mid[cellid] + n_verts = len(cell_xv) + vertex_points.append( + np.column_stack([cell_xv, cell_yv, np.full(n_verts, cell_z)]) + ) + vertex_cells.append(np.full(n_verts, cellid, dtype=int)) + + # Combine cell centers and vertices + all_points = np.vstack([centers] + vertex_points) + all_cells = np.concatenate([center_cells] + vertex_cells) + + points.extend(all_points.tolist()) + point_to_cell.extend(all_cells.tolist()) + + # Restore copy cache setting + self.grid._copy_cache = original_copy_cache + + def _precompute_hulls(self): + """Pre-compute ConvexHull equations and bounding boxes for all cells.""" + # Disable copy cache to avoid expensive deepcopy during precomputation + original_copy_cache = getattr(self.grid, "_copy_cache", True) + self.grid._copy_cache = False + + self.hull_equations = [] + self.bounding_boxes = np.zeros((self.ncells, 4)) # xmin, xmax, ymin, ymax + + for cellid in range(self.ncells): + verts = self._get_cell_vertices(cellid) + + # Handle empty or degenerate cells + if len(verts) < 3: + self.bounding_boxes[cellid] = [np.inf, -np.inf, np.inf, -np.inf] + self.hull_equations.append(None) + continue + + # Compute bounding box + self.bounding_boxes[cellid] = [ + verts[:, 0].min(), + verts[:, 0].max(), + verts[:, 1].min(), + verts[:, 1].max(), + ] + + # Compute ConvexHull equations + try: + hull = ConvexHull(verts) + self.hull_equations.append(hull.equations) + except Exception: + # Degenerate geometry + self.hull_equations.append(None) + + # Restore copy cache setting + self.grid._copy_cache = original_copy_cache + + def _precompute_3d_bounds(self): + """Pre-compute 3D bounding boxes (x,y,z bounds) for all cells.""" + xv, yv, zv = self.grid.xyzvertices + + # Store 3D bounding boxes: [xmin, xmax, ymin, ymax, zmin, zmax] + self.bounding_boxes_3d = np.zeros((self.ncells, 6)) + + for cellid in range(self.ncells): + cell_xv = np.atleast_1d(xv[cellid]) + cell_yv = np.atleast_1d(yv[cellid]) + cell_z_top = zv[0, cellid] + cell_z_bot = zv[1, cellid] + + self.bounding_boxes_3d[cellid] = [ + np.min(cell_xv), + np.max(cell_xv), + np.min(cell_yv), + np.max(cell_yv), + min(cell_z_top, cell_z_bot), + max(cell_z_top, cell_z_bot), + ] + + def _get_cell_vertices(self, cellid): + """Get vertices for a cell.""" + xv, yv, _ = self.grid.xyzvertices + return np.column_stack([xv[cellid], yv[cellid]]) + + def query_point(self, x, y, z=None, k=None): + """ + Find cell containing a single point. + + Parameters + ---------- + x, y : float + Point coordinates + z : float, optional + Z-coordinate. Required for 3D grids (grid_varies_by_layer=True). + k : int, optional + Number of unique cells to check (default 30) + + Returns + ------- + cellid : int or np.nan + Cell index containing the point, or np.nan if outside grid + """ + z_array = np.array([z]) if z is not None else None + result = self.query_points(np.array([x]), np.array([y]), z=z_array, k=k) + cellid = result[0] + return int(cellid) if not np.isnan(cellid) else np.nan + + def query_points(self, x, y, z=None, k=None): + """ + Find cells containing multiple points (vectorized). + + Uses KD-tree to find k nearest unique cells, then tests for containment. + For 2D grids: uses ConvexHull testing. + For 3D grids: uses 3D bounding box testing. + + Parameters + ---------- + x, y : array-like + Point coordinates (must have same length) + z : array-like, optional + Z-coordinates. Required for 3D grids (grid_varies_by_layer=True). + For 2D grids with layers, z-search is handled internally. + k : int, optional + Number of unique candidate cells to check per point (default 30) + + Returns + ------- + cellids : ndarray + Array of cell indices (np.nan for points outside grid). + For 3D grids, returns 3D cell index (nnodes). + For 2D grids, returns 2D cell index (ncpl). + """ + x = np.atleast_1d(x) + y = np.atleast_1d(y) + + if len(x) != len(y): + raise ValueError("x and y must have the same length") + + # For 3D grids, z is required + if self.is_3d: + if z is None: + raise ValueError( + "Z-coordinate required for 3D grids (grid_varies_by_layer=True)" + ) + z = np.atleast_1d(z) + if len(z) != len(x): + raise ValueError("z must have the same length as x and y") + else: + if z is not None: + z = np.atleast_1d(z) + if len(z) != len(x): + raise ValueError("z must have the same length as x and y") + + if k is None: + k = 30 # Default: check 30 unique cells + + # Build query points (2D or 3D) + if self.is_3d: + points = np.column_stack([x, y, z]) + else: + points = np.column_stack([x, y]) + + n_points = len(points) + results = np.full(n_points, np.nan, dtype=float) + + # For each point, query KD-tree to get k unique candidate cells + for i in range(n_points): + point = points[i] + + # Query KD-tree adaptively to get k unique cells + candidates = self._get_k_unique_cells(point, k) + + # Collect all matching cells for tie-breaking + matching_cells = [] + + # Test candidates in order (nearest first) + for cellid in candidates: + if self.is_3d: + # 3D: check 3D bounding box + if self._point_in_cell_3d(point, cellid): + matching_cells.append(cellid) + else: + # 2D: use ConvexHull test + if self._point_in_cell_vectorized(point[:2], cellid): + # Found 2D cell, now check z if provided + if z is None: + matching_cells.append(cellid) + else: + # Search through layers to find the right z + cell_3d = self._find_layer_for_z(cellid, z[i]) + if cell_3d is not None: + matching_cells.append(cell_3d) + + # Apply grid-specific tie-breaking when multiple cells match + if len(matching_cells) > 0: + results[i] = self._apply_tiebreaker(matching_cells) + + # Return int array if all points found, float array otherwise (to preserve nan) + valid_mask = ~np.isnan(results) + if np.all(valid_mask): + return results.astype(int) + return results + + def _get_k_unique_cells(self, point, k): + """ + Query KD-tree to get k unique candidate cells. + + Since KD-tree contains cell centers + vertices, we may need to query + more than k points to get k unique cells. + + Parameters + ---------- + point : ndarray + Query point [x, y] + k : int + Number of unique cells desired + + Returns + ------- + unique_cells : ndarray + Array of up to k unique cell indices, ordered by distance + """ + # Query k*5 points to get k unique cells (accounting for vertices) + k_points = min(k * 5, len(self.points)) + + distances, indices = self.tree.query(point, k=k_points) + + # Handle scalar result + if np.isscalar(indices): + indices = [indices] + + # Extract unique cells while preserving order + seen = set() + unique_cells = [] + for idx in indices: + cellid = self.point_to_cell[idx] + if cellid not in seen: + seen.add(cellid) + unique_cells.append(cellid) + if len(unique_cells) >= k: + break + + # If we still don't have k unique cells, query more points + while len(unique_cells) < k and k_points < len(self.points): + k_points = min(k_points * 2, len(self.points)) + distances, indices = self.tree.query(point, k=k_points) + + if np.isscalar(indices): + indices = [indices] + + for idx in indices: + cellid = self.point_to_cell[idx] + if cellid not in seen: + seen.add(cellid) + unique_cells.append(cellid) + if len(unique_cells) >= k: + break + + if k_points >= len(self.points): + break # Can't query any more points + + return np.array(unique_cells[:k], dtype=int) + + def _point_in_cell_vectorized(self, point, cellid): + """ + Test if point is inside cell using pre-computed bounding box + ConvexHull. + + Parameters + ---------- + point : ndarray + Point coordinates [x, y] + cellid : int + Cell index + + Returns + ------- + bool + True if point is inside cell + """ + # Fast bounding box rejection with epsilon tolerance for edge cases + # Expand bounding box slightly to handle points on boundaries + bbox = self.bounding_boxes[cellid] + if not ( + bbox[0] - self.epsilon <= point[0] <= bbox[1] + self.epsilon + and bbox[2] - self.epsilon <= point[1] <= bbox[3] + self.epsilon + ): + return False + + # Precise ConvexHull test + hull_eq = self.hull_equations[cellid] + if hull_eq is not None: + # Vectorized hyperplane test: Ax + By + C <= epsilon + # Tolerance allows points on edges to be included + distances = point @ hull_eq[:, :-1].T + hull_eq[:, -1] + return np.all(distances <= self.epsilon) + else: + # Degenerate geometry - should rarely happen + return False + + def _point_in_cell_3d(self, point, cellid): + """ + Test if 3D point is inside cell using 3D bounding box. + + Parameters + ---------- + point : ndarray + Point coordinates [x, y, z] + cellid : int + Cell index + + Returns + ------- + bool + True if point is inside cell's 3D bounding box + """ + # Use 3D bounding box test with epsilon tolerance + bbox = self.bounding_boxes_3d[cellid] + return ( + bbox[0] - self.epsilon <= point[0] <= bbox[1] + self.epsilon + and bbox[2] - self.epsilon <= point[1] <= bbox[3] + self.epsilon + and bbox[4] - self.epsilon <= point[2] <= bbox[5] + self.epsilon + ) + + def _apply_tiebreaker(self, matching_cells): + """ + Apply tie-breaking when multiple cells match. + + For vertex/unstructured grids: choose cell with lowest cell ID. + + Parameters + ---------- + matching_cells : list + List of cell indices that all contain the point + + Returns + ------- + cellid : int + The selected cell after tie-breaking + """ + if len(matching_cells) == 1: + return matching_cells[0] + + return min(matching_cells) + + def _find_layer_for_z(self, icell2d, z): + """ + Find 3D cell index for a 2D cell and z-coordinate. + + Searches through layers to find which layer contains the z-coordinate. + + Parameters + ---------- + icell2d : int + 2D cell index (from layer 0) + z : float + Z-coordinate + + Returns + ------- + cell_3d : int or None + 3D cell index, or None if z not found in any layer + """ + # Get z-bounds for all cells + _, _, zv = self.grid.xyzvertices + + # For VertexGrid: same 2D geometry for all layers + if self.grid.grid_type == "vertex": + # Search through all layers to find which contains z + # For VertexGrid: zv[lay, icell2d] = top, zv[lay+1, icell2d] = bottom + for lay in range(self.grid.nlay): + z_top = zv[lay, icell2d] + z_bot = zv[lay + 1, icell2d] + z_min = min(z_top, z_bot) + z_max = max(z_top, z_bot) + + if z_min <= z <= z_max: + # Found the layer! Return the layer index + return lay + + # z not found in any layer + return None + + # For UnstructuredGrid: compute 3D cell index using vectorized search + elif self.grid.grid_type == "unstructured": + # Build array of 3D cell indices for this 2D cell across all layers + if hasattr(self.grid, "ncpl") and isinstance(self.grid.ncpl, np.ndarray): + # Compute cumulative sum to get 3D cell indices + ncpl_cumsum = np.concatenate([[0], np.cumsum(self.grid.ncpl[:-1])]) + cell_indices_3d = icell2d + ncpl_cumsum + else: + # Fallback: assume constant ncpl + cell_indices_3d = ( + icell2d + np.arange(self.grid.nlay) * self.grid.ncpl[0] + ) + + # Get z-bounds for all layers of this 2D cell + z_tops = zv[0, cell_indices_3d] + z_bots = zv[1, cell_indices_3d] + z_mins = np.minimum(z_tops, z_bots) + z_maxs = np.maximum(z_tops, z_bots) + + # Find layers containing z (vectorized) + mask = (z_mins <= z) & (z <= z_maxs) + matching_layers = np.where(mask)[0] + + if len(matching_layers) > 0: + # Return first matching layer's 3D cell index + return cell_indices_3d[matching_layers[0]] + + return None + + def __repr__(self): + """String representation.""" + if self.grid.grid_type == "structured": + return ( + f"GeospatialIndex({self.grid.grid_type} grid, " + f"{self.grid.nrow} rows x {self.grid.ncol} cols)" + ) + return ( + f"GeospatialIndex({self.grid.grid_type} grid, " + f"{self.ncells} cells, " + f"{len(self.points)} indexed points)" + ) + + # ========================================================================= + # Unified intersect interface + # ========================================================================= + + def intersect(self, x, y, z=None, local=False, forgive=False): + """ + Find the cell(s) containing the given point(s). + + This is the unified interface for spatial queries across all grid types. + Dispatches to the optimal algorithm based on grid type: + - StructuredGrid: searchsorted (O(log n)) + - VertexGrid/UnstructuredGrid: KD-tree + ConvexHull + + Parameters + ---------- + x : float or array-like + The x-coordinate(s) of the query point(s) + y : float or array-like + The y-coordinate(s) of the query point(s) + z : float, array-like, or None + Optional z-coordinate(s). If provided, returns layer information. + local : bool, optional + If True, x and y are in local coordinates (default False) + forgive : bool, optional + If True, return NaN for points outside the grid instead of + raising an error (default False) + + Returns + ------- + For StructuredGrid: + row, col : int or ndarray + Row and column indices. If z is provided, returns (lay, row, col). + For VertexGrid: + cellid : int or ndarray + Cell index (icell2d). If z is provided, returns (lay, cellid). + For UnstructuredGrid: + cellid : int or ndarray + Cell index. If z is provided and grid_varies_by_layer is False, + returns (lay, cellid). + + Raises + ------ + ValueError + If point is outside grid and forgive=False + """ + if self.grid.grid_type == "structured": + return self._intersect_structured(x, y, z, local, forgive) + else: + return self._intersect_unstructured(x, y, z, local, forgive) + + def _intersect_structured(self, x, y, z=None, local=False, forgive=False): + """ + Find row/col for structured grid using searchsorted. + + Uses vectorized binary search for O(log n) performance. + """ + # 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 to local coordinates if needed + if not local: + x, y = self.grid.get_local_coords(x, y) + + # Get cached edge arrays + xe = self._xe + ye_flipped = self._ye_flipped + + # Vectorized row/col calculation + n_points = len(x) + rows = np.full(n_points, np.nan, dtype=float) + cols = np.full(n_points, np.nan, dtype=float) + + # Vectorized column finding using searchsorted + # side="left" ensures x==edge goes to the cell on the left (tie-breaking) + cols_valid = np.searchsorted(xe, x, side="left") - 1 + cols_mask = (cols_valid >= 0) & (cols_valid < self.grid.ncol) + cols[cols_mask] = cols_valid[cols_mask] + + # Vectorized row finding using searchsorted on flipped ye + # side="right" on flipped array ensures y==edge goes to lower row + rows_flipped = np.searchsorted(ye_flipped, y, side="right") + rows_valid = len(self._ye) - 1 - rows_flipped + rows_mask = (rows_valid >= 0) & (rows_valid < self.grid.nrow) + rows[rows_mask] = rows_valid[rows_mask] + + # Check for errors if not forgiving + if not forgive: + invalid_mask = np.isnan(rows) | np.isnan(cols) + if np.any(invalid_mask): + idx = np.where(invalid_mask)[0][0] + raise ValueError( + f"x, y point given is outside of the model area: " + f"({x[idx]}, {y[idx]})" + ) + + # 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 2D 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 - vectorized layer finding + lays = np.full(n_points, np.nan, dtype=float) + + # Only process points that have valid row/col + valid_mask = ~(np.isnan(rows) | np.isnan(cols)) + valid_indices = np.where(valid_mask)[0] + + if len(valid_indices) > 0: + valid_rows = rows[valid_indices].astype(int) + valid_cols = cols[valid_indices].astype(int) + valid_z = z[valid_indices] + + # Get top/bottom elevations for all valid points and all layers + tops_bottoms = self.grid.top_botm[:, valid_rows, valid_cols].T + + # Check which layer each point belongs to + in_layer = (tops_bottoms[:, :-1] >= valid_z[:, np.newaxis]) & ( + valid_z[:, np.newaxis] >= tops_bottoms[:, 1:] + ) + + # Find the first (topmost) layer for each point + layer_indices = np.argmax(in_layer, axis=1) + + # Set layer values only where a valid layer was found + n_valid = len(valid_indices) + found_layer = in_layer[np.arange(n_valid), layer_indices] + lays[valid_indices[found_layer]] = layer_indices[found_layer] + + # Check for errors if not forgiving + if not forgive: + not_found = ~found_layer + if np.any(not_found): + idx = valid_indices[not_found][0] + raise ValueError( + f"point given is outside the model area: " + f"({x[idx]}, {y[idx]}, {z[idx]})" + ) + + # Return 3D 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) + 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 _intersect_unstructured(self, x, y, z=None, local=False, forgive=False): + """ + Find cell(s) for vertex/unstructured grid using KD-tree. + + Uses KD-tree nearest neighbor search + ConvexHull point-in-polygon. + """ + # 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 + 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 to world coordinates if local + if local: + x, y = self.grid.get_coords(x, y) + + # Use existing query_points for the spatial search + if self.is_3d: + # 3D grid requires z + if z is None: + raise ValueError( + "Z-coordinate required for 3D grids (grid_varies_by_layer=True)" + ) + cellids = self.query_points(x, y, z=z) + else: + # 2D search, then layer search if z provided + cellids = self.query_points(x, y, z=z) + + # Check for errors if not forgiving + if not forgive: + invalid_mask = np.isnan(cellids) + if np.any(invalid_mask): + idx = np.where(invalid_mask)[0][0] + if z is not None: + raise ValueError( + f"point given is outside the model area: " + f"({x[idx]}, {y[idx]}, {z[idx]})" + ) + else: + raise ValueError( + f"x, y point given is outside of the model area: " + f"({x[idx]}, {y[idx]})" + ) + + # Handle return format based on grid type and z + if self.grid.grid_type == "vertex" and z is not None and not self.is_3d: + # For VertexGrid with z, return (lay, icell2d) + # The layer was already found in query_points via _find_layer_for_z + # We need to extract layer and icell2d from the result + n_points = len(x) + lays = np.full(n_points, np.nan, dtype=float) + icell2ds = np.full(n_points, np.nan, dtype=float) + + valid_mask = ~np.isnan(cellids) + if np.any(valid_mask): + # Re-query 2D to get icell2d, then find layer separately + for i in np.where(valid_mask)[0]: + # Re-do the 2D query to get icell2d + icell2d_result = self.query_points( + np.array([x[i]]), np.array([y[i]]) + )[0] + if not np.isnan(icell2d_result): + icell2ds[i] = icell2d_result + # Find layer for this cell + lay = self._find_layer_for_z(int(icell2d_result), z[i]) + if lay is not None: + lays[i] = lay + + if is_scalar_input: + lay, icell2d = lays[0], icell2ds[0] + if not np.isnan(lay) and not np.isnan(icell2d): + return int(lay), int(icell2d) + return lay, icell2d + else: + valid_3d = ~np.isnan(lays) & ~np.isnan(icell2ds) + return ( + lays.astype(int) if np.all(valid_3d) else lays, + icell2ds.astype(int) if np.all(valid_3d) else icell2ds, + ) + + # Simple case: return cellid(s) directly + if is_scalar_input: + cellid = cellids[0] + if not np.isnan(cellid): + return int(cellid) + return cellid + else: + valid_mask = ~np.isnan(cellids) + if np.all(valid_mask): + return cellids.astype(int) + return cellids