perf(geospatial): optimize grid.intersect() with KD-tree + ConvexHull

[aacovski]

Summary

Optimize grid.intersect() performance across all grid types using KD-tree spatial indexing with pre-computed ConvexHull equations and fully vectorized structured grid with searchsorted. This is all on the assumption that grids will never be non-convex and always layered.

Performance Improvements

geospatial_performance_results.csv
~50-300x speedup

Key Optimizations

1. KD-tree + vertices indexing: O(log n) candidate cell search using centroids + vertices
2. Pre-computed ConvexHull equations: Vectorized point-in-polygon testing via hyperplane equations
3. Bounding box rejection: Fast elimination before expensive polygon tests
4. Deepcopy fix: Disabled _copy_cache during index build to avoid bottleneck
5. Vectorized result processing: Eliminated Python loops in grid.intersect() methods

Changes

• Add flopy/utils/geospatial_index.py with GeospatialIndex class
• Optimize VertexGrid.intersect() - 2D kdtree, 2D convexhull and layer search
• Optimize UnstructuredGrid.intersect() varies_by_layers=False - 2D kdtree, 2D convexhull and layer search
• Optimize UnstructuredGrid.intersect() varies_by_layers=True - 3D kdtree, 3D convexhull
• Optimize StructuredGrid.intersect() with vectorized row/col finding - searchsorted
• Add edge case tests for thin sliver cells, boundaries, corners
• GeospatialIndex now raises ValueError if passed a StructuredGrid
• All methods now return np.nan for outside points (previously StructuredGrid returned (None, nan, nan) with z, and GeospatialIndex.query_point returned None)
• Array queries return int64 when all points inside, float64 when nan values present
• Matches behavior of native VertexGrid.intersect() and UnstructuredGrid.intersect() on develop
• 1:1 point-to-cell mapping
• All methods guarantee same number of outputs as inputs
• Outside points get nan instead of being dropped or raising errors (with forgive=True)

Test coverage

• ~38 tests covering return type consistency, tie-breaking, 3D layered grids, and boundary handling

Regressions

• grids with less than 10 cells or points perform worse due to overhead associated with new methods, but execute in less than 0.5 ms so it doesn't matter.

[codecov]

Codecov Report

❌ Patch coverage is 68.70748% with 92 lines in your changes missing coverage. Please review.
✅ Project coverage is 72.6%. Comparing base (556c088) to head (f1996f6).
⚠️ Report is 87 commits behind head on develop.

Files with missing lines	Patch %	Lines
flopy/utils/geospatial_index.py	65.8%	80 Missing ⚠️
flopy/discretization/structuredgrid.py	87.8%	4 Missing ⚠️
flopy/discretization/unstructuredgrid.py	66.6%	4 Missing ⚠️
flopy/discretization/vertexgrid.py	73.3%	4 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##           develop    #2654      +/-   ##
===========================================
+ Coverage     55.5%    72.6%   +17.1%     
===========================================
  Files          644      669      +25     
  Lines       124135   129738    +5603     
===========================================
+ Hits         68947    94263   +25316     
+ Misses       55188    35475   -19713     

Files with missing lines	Coverage Δ
flopy/discretization/structuredgrid.py	53.2% <87.8%> (+5.8%)	⬆️
flopy/discretization/unstructuredgrid.py	75.7% <66.6%> (-5.7%)	⬇️
flopy/discretization/vertexgrid.py	79.2% <73.3%> (-4.4%)	⬇️
flopy/utils/geospatial_index.py	65.8% <65.8%> (ø)

... and 557 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[aacovski]

@wpbonelli I’m looking for some guidance on how to handle these tests. The issue arises when a point lies exactly on the vertices of grid cells, making it ambiguous which of the 3-n cells it should belong to:

FAILED test_grid.py::test_intersection - assert 268 == 267

In these non-deterministic cases, different methods return different cell indices. I can force the kdtree-based methods to return the expected value, but that comes with a performance cost. It then causes another test failure (assert DIS == DISV), since those methods also behave differently. I could also update the structured grid path to enforce consistent behavior so all tests pass.

Before making those changes, I’d like your input on the preferred approach to resolve these cascading test failures.

[aacovski]

No performance hit with changing tiebreaking

[dbrakenhoff]

@wpbonelli I’m looking for some guidance on how to handle these tests. The issue arises when a point lies exactly on the vertices of grid cells, making it ambiguous which of the 3-n cells it should belong to:

@aacovski in the gridintersect module, in case of ties, we decided to always return the lowest grid cell number. But I think that any approach as long as it is clearly documented (and maybe consistent across runs) should be fine?

[wpbonelli]

Agree with @dbrakenhoff I think lowest number is the existing convention so might as well stick to it if there is no performance cost?

[aacovski]

Great, that's what I did. I'll consider adding more tests for coverage, but otherwise I think the PR is ready. Thanks

[wpbonelli]

across all grid types

grids will never be non-convex and always layered

does the latter mean the new grid index and related performance boost don't support DISU?

[aacovski]

does the latter mean the new grid index and related performance boost don't support DISU?

As I understand it, “fully 3D unstructured” grids are more characteristic of FEM models, and I’m not aware of MODFLOW supporting that kind of discretization. That said, my code does support DISU in the sense that varies_by_layers can be set to either True or False... I'm not aware of any FDM or FEM software supporting non-convex elemnts/grids. Please let me know if I'm missing anything.

[aacovski]

I am going to clean up docstrings later today as I noticed some old info. Sorry about that. I want to add some more tests too.

[wpbonelli]

@aacovski I guess I just didn't properly understand what you meant by "always layered"

[wpbonelli]

it's maybe worth noting that the index assumes that x/y cell centers are in fact centroid coords? presumably the index requires this, rather than simple geometric mean, for polygons for which they aren't equivalent? maybe we should clarify in the docstrings on the grid classes as "cell centers" is ambiguous?

[aacovski]

it's maybe worth noting that the index assumes that x/y cell centers are in fact centroid coords?

I hadn’t considered that the provided cell centers might not correspond to geometric centroids even for convex polygons. The index shouldn’t assume that, so I’ll update to clarify how cell centers are interpreted. Thanks for pointing it out!

[wpbonelli]

The index shouldn’t assume that [the provided cell centers correspond to geometric centroids], so I’ll update to clarify how cell centers are interpreted.

Ah, ok- and the index doesn't even assume the center is inside the cell? It just performs better if so?

I hadn’t considered that the provided cell centers might not correspond to geometric centroids even for convex polygons.

I think the intention was/is for it to be cell centroid, even though we don't check this when provided by the user or read from a binary grid file. Evidence:

• flopy/flopy/mf6/utils/binarygrid_util.py

  Line 327 in e42f511

  def _get_cellcenters(self):

  defines cell centers as centroids, maybe (but in any case suggestively) as a typo

  • flopy/flopy/mf6/utils/binarygrid_util.py

    Line 329 in e42f511

    Get the cell centers centroids for a MODFLOW 6 GWF model that uses

• flopy/flopy/utils/cvfdutil.py

  Line 17 in e42f511

  def centroid_of_polygon(points):

  is used when calculating x/y cell centers in other routines

  • flopy/flopy/utils/cvfdutil.py

    Line 143 in e42f511

    def to_cvfd(

  • flopy/flopy/utils/cvfdutil.py

    Line 306 in e42f511

    def shapefile_to_xcyc(shp):

  • flopy/flopy/utils/cvfdutil.py

    Line 369 in e42f511

    def get_disv_gridprops(verts, iverts, xcyc=None):

• flopy/flopy/utils/geometry.py

  Line 739 in e42f511

  def get_polygon_centroid(geom):

  is used from

  • flopy/flopy/discretization/unstructuredgrid.py

    Line 1062 in e42f511

    def from_argus_export(cls, file_path, nlay=1):

Maybe something to clarify in the MF6IO guide. For DIS it doesn't matter but it becomes relevant for DISV/U.

[wpbonelli]

@aacovski it wasn't really clear from previous comment but I think it's best to revert the docstrings to say centroid rather than leaving it ambiguous. my first comment was really to point out the (unwanted?) ambiguity in the existing grid classes, not to say your new ones should change

[wpbonelli]

Another thought: why not move the structured grid searchsorted query implementation to GeospatialIndex? That way all the Grid.intersect methods could use the index. Since the index receives the grid on init, it knows what kind it is. This seems a bit neater, and also opens the way for cases where a structured grid query might still find the KD tree useful, e.g. ball queries.

Also just want to say, thanks for working on this. This seems like it will be a great asset both in the near term and longer term- in 4.x it could sit behind xarray-based user APIs, whether builtin .sel syntax (point queries) or, for the fancy stuff, maybe an accessor.

[aacovski]

I'll work on implementing that: moving intersect to base class and remove the subclass intersect methods. Regarding the docstrings: since there is a difference between centroid and centers, based on everything I've read, it makes more sense to use the more general case (centers) rather than centroids which could be inaccurate.

[wpbonelli]

I'm pretty sure "cell centers" is meant as centroids, that's why I'm suggesting to change the doctrings up in the grids themselves. But I could be wrong. Hopefully someone with more tenure can chime in.