feat(geospatial): add 3D spatial indexing and vectorize z-coordinate search

Add comprehensive 3D support to GeospatialIndex with automatic detection
and two-pronged approach for handling grids with varying layer geometry.

**3D Spatial Indexing (grid_varies_by_layer=True):**
- Uses (x,y,z) KD-tree for grids where geometry varies by layer
- Implements 3D bounding box containment tests
- Indexes all nnodes with full 3D coordinates
- Performance: up to 414x speedup for 15000 cells with 10000 points

**2D with Z-Search (grid_varies_by_layer=False):**
- Uses (x,y) KD-tree with vectorized z-coordinate layer search
- Suitable for VertexGrid and UnstructuredGrid with consistent 2D geometry
- Performance: up to 348x speedup for 5000 cells with 10000 points

**Key Changes:**
- Add is_3d flag auto-detected from grid.grid_varies_by_layer
- Implement _build_3d_index() for 3D KD-tree construction
- Implement _precompute_3d_bounds() for 3D bounding box tests
- Implement _point_in_cell_3d() for 3D containment
- Vectorize _find_layer_for_z() using numpy masks (no loops)
- Update query_points() to accept z parameter and route to 2D/3D logic
- Simplify UnstructuredGrid.intersect() and VertexGrid.intersect()
- Fix list vs array handling with np.atleast_1d()

**Performance Benchmarks:**
Structured grids: 0.10x-0.58x (slower, should use searchsorted)
Vertex grids: 1.2x-348x speedup (best for 1000+ cells)
Unstructured 2D: 1.1x-352x speedup (best for 1000+ cells)
Unstructured 3D: 1.1x-414x speedup (best for 2000+ cells)

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <[email protected]>

Author: adacovsk

flopy/discretization/unstructuredgrid.py

@@ -775,53 +775,28 @@ def intersect(self, x, y, z=None, local=False, forgive=False):
         if not hasattr(self, "_geospatial_index") or self._geospatial_index is None:
             self._geospatial_index = GeospatialIndex(self)
 
-        # Use KD-tree for fast spatial queries
-        n_points = len(x)
-        cellids = self._geospatial_index.query_points(x, y)
+        # Use GeospatialIndex for spatial queries
+        # For grid_varies_by_layer=True, z is required (3D query)
+        # For grid_varies_by_layer=False, z-search handled internally
+        cellids = self._geospatial_index.query_points(x, y, z=z)
 
-        xv, yv, zv = self.xyzvertices
-
-        # Initialize result array - cellids already contains np.nan for not found
+        # Initialize result array
         results = cellids.copy()
 
-        # Vectorized error checking for points not found
+        # Error checking for points not found
         if not forgive:
             unfound_mask = np.isnan(cellids)
             if np.any(unfound_mask):
                 idx = np.where(unfound_mask)[0][0]
-                xi, yi = x[idx], y[idx]
-                raise ValueError(f"point ({xi}, {yi}) is outside of the model area")
-
-        # Only process z-coordinates if provided
-        if z is not None:
-            # Only loop over valid points (that were found in 2D)
-            valid_mask = ~np.isnan(cellids)
-            valid_indices = np.where(valid_mask)[0]
-
-            for i in valid_indices:
-                icell2d = int(cellids[i])
-                zi = z[i]
-                found = False
-
-                # 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 not found:
-                    if forgive:
-                        results[i] = np.nan
-                    else:
-                        xi, yi = x[i], y[i]
-                        raise ValueError(
-                            f"point ({xi}, {yi}, {zi}) is outside of the model area"
-                        )
+                if z is not None:
+                    raise ValueError(
+                        f"point ({x[idx]}, {y[idx]}, {z[idx]}) is outside of "
+                        f"the model area"
+                    )
+                else:
+                    raise ValueError(
+                        f"point ({x[idx]}, {y[idx]}) is outside of the model area"
+                    )
 
         # Return scalar if input was scalar, otherwise return array
         if is_scalar_input:

flopy/discretization/vertexgrid.py

@@ -405,43 +405,36 @@ def intersect(self, x, y, z=None, local=False, forgive=False):
         if not hasattr(self, "_geospatial_index") or self._geospatial_index is None:
             self._geospatial_index = GeospatialIndex(self)
 
-        # Use KD-tree for fast spatial queries
-        n_points = len(x)
-        cellids = self._geospatial_index.query_points(x, y)
+        # Use GeospatialIndex for spatial queries
+        # For VertexGrid, z-search is handled internally by GeospatialIndex
+        # which returns layer index
+        if z is not None:
+            lays_raw = self._geospatial_index.query_points(x, y, z=z)
+            # GeospatialIndex returns layer index for VertexGrid
+            lays = lays_raw.copy()
+            # Get icell2d from 2D query (all layers have same 2D geometry)
+            cellids = self._geospatial_index.query_points(x, y, z=None)
+        else:
+            cellids = self._geospatial_index.query_points(x, y, z=None)
 
         # Vectorized processing of results
-        # cellids is already np.nan for not found, convert to the right type
         results = cellids.copy()
 
         # Check for unfound points if not forgiving
         if not forgive:
             unfound_mask = np.isnan(cellids)
             if np.any(unfound_mask):
                 idx = np.where(unfound_mask)[0][0]
-                raise ValueError(
-                    f"point given is outside of the model area: ({x[idx]}, {y[idx]})"
-                )
-
-        # Find z-layers if z is provided (vectorized where possible)
-        if z is not None:
-            lays = np.full(n_points, np.nan, dtype=float)
-
-            # Only process points that were found
-            valid_mask = ~np.isnan(cellids)
-            valid_indices = np.where(valid_mask)[0]
-
-            for i in valid_indices:
-                icell2d = int(cellids[i])
-                zi = z[i]
-                # Find layer for this point
-                for lay in range(self.nlay):
-                    if (
-                        self.top_botm[lay, icell2d]
-                        >= zi
-                        >= self.top_botm[lay + 1, icell2d]
-                    ):
-                        lays[i] = lay
-                        break
+                if z is not None:
+                    raise ValueError(
+                        f"point given is outside of the model area: "
+                        f"({x[idx]}, {y[idx]}, {z[idx]})"
+                    )
+                else:
+                    raise ValueError(
+                        f"point given is outside of the model area: "
+                        f"({x[idx]}, {y[idx]})"
+                    )
 
         # Return results
         if z is None: