fix(geospatial): improve ncpl handling and use cKDTree for performance

This commit improves the GeospatialIndex class to better handle different
grid types:

- Use cKDTree instead of KDTree for better performance in spatial queries
- Add proper handling for VertexGrid (scalar ncpl) vs UnstructuredGrid
  (array ncpl)
- Handle grid_varies_by_layer case for unstructured grids where each layer
  has different geometry
- Fix line length issues in error messages for code style compliance

Note: There is a known test failure in test_intersection where structured
grids using searchsorted() return different results than DISV grids using
the KD-tree approach for edge cases. This will be addressed in a separate
commit.

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

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

Author: adacovsk

flopy/discretization/structuredgrid.py

@@ -1050,7 +1050,8 @@ def intersect(self, x, y, z=None, local=False, forgive=False):
             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: ({x[idx]}, {y[idx]})"
+                    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

flopy/discretization/unstructuredgrid.py

@@ -819,7 +819,10 @@ def intersect(self, x, y, z=None, local=False, forgive=False):
                         results[i] = np.nan
                     else:
                         xi, yi = x[i], y[i]
-                        raise ValueError(f"point ({xi}, {yi}, {zi}) is outside of the model area")
+                        raise ValueError(
+                            f"point ({xi}, {yi}, {zi}) is outside of the "
+                            f"model area"
+                        )
 
         # Return scalar if input was scalar, otherwise return array
         if is_scalar_input:

flopy/utils/geospatial_index.py

@@ -7,7 +7,7 @@
 """
 
 import numpy as np
-from scipy.spatial import KDTree, ConvexHull
+from scipy.spatial import ConvexHull, cKDTree
 
 
 class GeospatialIndex:
@@ -69,10 +69,8 @@ def __init__(self, grid):
 
     def _build_index(self):
         """
-        Build KD-tree with centroids + vertices and pre-compute ConvexHull equations.
-
-        Strategy: Include both centroids and vertices in KD-tree to handle edge cases,
-        but use smart querying to ensure we get k unique cells.
+        Build KD-tree with centroids + vertices and pre-compute
+        ConvexHull equations.
         """
         points = []
         point_to_cell = []
@@ -83,7 +81,22 @@ def _build_index(self):
             self._build_structured_index(points, point_to_cell)
         elif self.grid.grid_type in ('vertex', 'unstructured'):
             if hasattr(self.grid, 'ncpl'):
-                self.ncells = self.grid.ncpl
+                ncpl = self.grid.ncpl
+                if np.isscalar(ncpl):
+                    # VertexGrid: ncpl is scalar
+                    self.ncells = ncpl
+                else:
+                    # UnstructuredGrid: ncpl is array
+                    if (hasattr(self.grid, 'grid_varies_by_layer')
+                            and self.grid.grid_varies_by_layer):
+                        # Each layer has different geometry
+                        self.ncells = self.grid.nnodes
+                    else:
+                        # Same 2D grid for all layers
+                        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)
@@ -93,7 +106,7 @@ def _build_index(self):
         # Build KD-tree from centroids + vertices
         self.points = np.array(points)
         self.point_to_cell = np.array(point_to_cell, dtype=int)
-        self.tree = KDTree(self.points)
+        self.tree = cKDTree(self.points)
 
         # Pre-compute ConvexHull equations and bounding boxes
         self._precompute_hulls()
@@ -144,11 +157,7 @@ def _build_vertex_index(self, points, point_to_cell):
         self.grid._copy_cache = original_copy_cache
 
     def _precompute_hulls(self):
-        """
-        Pre-compute ConvexHull equations and bounding boxes for all cells.
-
-        This is the key optimization - compute once, reuse many times.
-        """
+        """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
@@ -291,8 +300,7 @@ def _get_k_unique_cells(self, point, k):
         unique_cells : ndarray
             Array of up to k unique cell indices, ordered by distance
         """
-        # Start by querying k*5 points (centroid + ~4 vertices per cell)
-        # This should give us at least k unique cells in most cases
+        # 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)