Suggestion: vtk/pyvista for spatial indexing

[Huite]

Just a suggestion: I spent a fair amount of time setting up the spatial index and querying in https://github.com/Deltares/numba_celltree for specific logic (overlaps, intersections, etc.).
VTK and pyvista might be a worthwhile alternative:

# Each cell begins with the number of points in the cell and then the points
# composing the cell
counts = (face_node_connectivity != FILL).sum(axis=1)
cells2d = np.column_stack((counts, grid.face_node_connectivity))

mesh = pv.UnstructuredGrid(
    cells2d[cells2d != -1],
    np.full(grid.n_face, CellType.TRIANGLE),
    np.column_stack((node_x, node_y, np.zeros(grid.n_node))),
)

indices = mesh.find_containing_cell(np.array([[150_000.0, 463_000.0, 0.0]]))

Essentially all the queries you'd like are available:

mesh.find_cells_within_bounds
mesh.find_cells_along_line
mesh.find_cells_within_bounds
mesh.find_closest_cell
mesh.find_closest_point
mesh.find_containing_cell

It might not solve everything since I'm not sure yet what VTK supports in terms of non-Cartesian coordinates, but this makes 3D topologies very straightforward to support as well!

[Huite]

Bit more digging: performance might be a concern.
pyvista accepts numpy arrays, but this mostly calls Python for loops to unvectorized vtk methods, which is a bit disappointing.

Ideally, you'd use compiled parallel for loops (like Cython or Numba prange).

I think pyvista also reinitialized the locator every query, but this can be avoided by using the vtk locator directly.

[philipc2]

Hi @Huite

Thanks for sharing this! This does seem very relevant to a lot of the work that we've been putting together.

Although, I am a little hesitant in how applicable this would be to our specific use case of spherical unstructured grids. We've been working with @hongyuchen1030 to implement many of the routines above directly on the unit sphere.

We achieve mesh.find_cells_within_bounds and mesh.find_cells_along_line by pre-computing a spherical bounding box for each of our faces and use that for our queries. There is some overhead in initially computing the bounding box, but we've drastically reduced it by using Numba.

The bounds are currently stored as an array, however I could imagine there are more efficient ways to organize the bounds that would reduce our query time (@hongyuchen1030 if I remember correctly, I think this is something you presented or showcased before

@aaronzedwick is currently working on a Point in Polygon implementation in #1056

We have an approximate version of mesh.find_closest_cell and mesh.find_closest_point that works well-enough built around a BallTree, though this is something we'd like to improve moving forward with a more robust implementation.

Ideally, you'd use compiled parallel for loops (like Cython or Numba prange).

This is what we've been leaning towards (Numba) for some time now, and it's been very successful.

I spent a fair amount of time setting up the spatial index and querying in https://github.com/Deltares/numba_celltree

I've glanced over this before and the implementation looks great. I could see this being useful, though my same concerns above regarding spherical geometry apply.

We could however possibly utilize these packages if we chose to project our grid to 2D using something like a sterographic projection, though that also introduces some new difficulties.

[Huite]

Ah I see, I had seen some methods for spherical geometries (the Gaussian quadrature methods, etc.), but hadn't realized the explicit focus so much.

I suppose you have come across this?
https://github.com/benbovy/spherely

This spatial index might also be relevant, if I understand correctly: https://github.com/benbovy/pys2index

In terms of inspiration, R has the simple features package which seems to enable both GEOS (Cartesian) and S2geometry (spherical). It's obviously much more akin to geopandas in terms of functionality and setup, but there's almost certainly interesting things in there, see: https://r-spatial.github.io/sf/articles/sf7.html