feat(GridIntersect): add array support, fast point locating and support for points with z-coordinates

.docs/Notebooks/grid_intersection_example.py

@@ -4,17 +4,28 @@
 #     notebook_metadata_filter: all
 #     text_representation:
 #       extension: .py
-#       format_name: light
-#       format_version: '1.5'
-#       jupytext_version: 1.14.5
+#       format_name: percent
+#       format_version: '1.3'
+#       jupytext_version: 1.18.1
 #   kernelspec:
-#     display_name: Python 3 (ipykernel)
+#     display_name: flopy
 #     language: python
 #     name: python3
+#   language_info:
+#     codemirror_mode:
+#       name: ipython
+#       version: 3
+#     file_extension: .py
+#     mimetype: text/x-python
+#     name: python
+#     nbconvert_exporter: python
+#     pygments_lexer: ipython3
+#     version: 3.13.7
 #   metadata:
 #     section: dis
 # ---
 
+# %% [markdown]
 # # <a id="top"></a>Intersecting model grids with shapes
 #
 # _Note: This feature requires the shapely package (which is an optional FloPy dependency)._
@@ -35,9 +46,10 @@
 #     - [MultiLineString with triangular grid](#trigrid.2)
 #     - [MultiPoint with triangular grid](#trigrid.3)
 
+# %% [markdown]
 # Import packages
 
-# +
+# %%
 import sys
 
 import matplotlib as mpl
@@ -56,8 +68,8 @@
 print(f"matplotlib version: {mpl.__version__}")
 print(f"shapely version: {shapely.__version__}")
 print(f"flopy version: {flopy.__version__}")
-# -
 
+# %% [markdown]
 # ## <a id="gridclass"></a>[GridIntersect Class](#top)
 #
 # The GridIntersect class is constructed by passing a flopy modelgrid object to
@@ -67,7 +79,7 @@
 # - `rtree`: either `True` (default) or `False`. When True, an STR-tree is built,
 # which allows for fast spatial queries. Building the STR-tree takes some
 # time however. Setting the option to False avoids building the STR-tree but requires
-# the intersection calculation to loop through all grid cells. It is generally
+# the intersection calculation to loop through more candidate grid cells. It is generally
 # recommended to set this option to True.
 # - `local`: either `False` (default) or `True`. When True the local model coordinates
 # are used. When False, real-world coordinates are used. Can be useful if shapes are
@@ -76,73 +88,98 @@
 # The important methods in the GridIntersect object are:
 #
 # - `intersects()`: returns cellids for gridcells that intersect a shape (accepts
-# shapely geometry objects, flopy geometry object, shapefile.Shape objects, and
-# geojson objects)
+# shapely geometry objects, arrays of shapely geometry object, flopy geometry object,
+# shapefile.Shape objects, and geojson objects).
 # - `intersect()`: for intersecting the modelgrid with point, linestrings, and
 # polygon geometries (accepts shapely geometry objects, flopy geometry object,
 # shapefile.Shape objects, and geojson objects)
+# - `points_to_cellids()`: for quickly locating points in the modelgrid and
+# getting a 1:1 mapping of points to cellids. Especially useful when passing in array
+# of points.
 # - `ix.plot_intersection_result()`: for plotting intersection results
 #
 # In the following sections examples of intersections are shown for structured
 # and vertex grids for different types of shapes (Polygon, LineString and Point).
 
+# %% [markdown]
 # ## <a id="rectgrid"></a>[Rectangular regular grid](#top)
 
+# %%
 delc = 10 * np.ones(10, dtype=float)
 delr = 10 * np.ones(10, dtype=float)
 
+# %%
 xoff = 0.0
 yoff = 0.0
 angrot = 0.0
 sgr = fgrid.StructuredGrid(
     delc, delr, top=None, botm=None, xoff=xoff, yoff=yoff, angrot=angrot
 )
 
+# %%
 sgr.plot()
-
+# %% [markdown]
 # ### <a id="rectgrid.1"></a>[Polygon with regular grid](#top)
 # Polygon to intersect with:
 
+# %%
 p = Polygon(
     shell=[(15, 15), (20, 50), (35, 80.0), (80, 50), (80, 40), (40, 5), (15, 12)],
     holes=[[(25, 25), (25, 45), (45, 45), (45, 25)]],
 )
 
+# %%
 # Create the GridIntersect class for our modelgrid.
-# TODO: remove method kwarg in v3.9.0
-ix = GridIntersect(sgr, method="vertex")
+ix = GridIntersect(sgr)
 
+# %% [markdown]
 # Do the intersect operation for a polygon
 
+# %%
 result = ix.intersect(p)
 
+# %% [markdown]
 # The results are returned as a numpy.recarray containing several fields based on the intersection performed. An explanation of the data in each of the possible fields is given below:
 # - **cellids**: contains the cell ids of the intersected grid cells
-# - **vertices**: contains the vertices of the intersected shape
 # - **areas**: contains the area of the polygon in that grid cell (only for polygons)
 # - **lengths**: contains the length of the linestring in that grid cell (only for linestrings)
 # - **ixshapes**: contains the shapely object representing the intersected shape (useful for plotting the result)
 #
-# Looking at the first few entries of the results of the polygon intersection (convert to pandas.DataFrame for prettier formatting)
+# Looking at the first few entries of the results of the polygon intersection. Note that you can convert the result to a GeoDataFrame (if geopandas is installed) with `geo_dataframe=True`.
 
-result[:5]
-# pd.DataFrame(result).head()
+# %%
+ix.intersect(p, geo_dataframe=True).head()
 
+# %% [markdown]
 # The cellids can be easily obtained
 
+# %%
 result.cellids
 
+# %% [markdown]
 # Or the areas
 
+# %%
 result.areas
 
-# If  a user is only interested in which cells the shape intersects (and not the areas or the actual shape of the intersected object) with there is also the `intersects()` method. This method works for all types of shapely geometries.
+# %% [markdown]
+# If a user is only interested in which cells the shape intersects (and not the areas or
+# the actual shape of the intersected object) with there is also the `intersects()`
+# method. This method works for all types of shapely geometries including arrays of
+# shapely geometries.
+#
+# This method returns `shp_ids` and `cellids` fields. The `shp_ids` field contains the
+# index of the geometry in the original input shape provided by the user. This is useful
+# when the input is an array of shapely geometries. In this case we have only one polygon,
+# so the `shp_id` is always equal to 0.
 
+# %%
 ix.intersects(p)
 
+# %% [markdown]
 # The results of an intersection can be visualized with the `GridIntersect.plot_intersection_result()` method.
 
-# +
+# %%
 # create a figure and plot the grid
 fig, ax = plt.subplots(1, 1, figsize=(8, 8))
 
@@ -160,8 +197,7 @@
 
 # add legend
 ax.legend([h2], [i.get_label() for i in [h2]], loc="best")
-# -
-
+# %% [markdown]
 # The `intersect()` method contains several keyword arguments that specifically deal with polygons:
 #
 # - `contains_centroid`: only store intersection result if cell centroid is contained within polygon
@@ -171,7 +207,7 @@
 #
 # Example with `contains_centroid` set to True, only cells in which centroid is within the intersected polygon are stored. Note the difference with the previous result.
 
-# +
+# %%
 # contains_centroid example
 
 result2 = ix.intersect(p, contains_centroid=True)
@@ -190,11 +226,10 @@
 
 # add legend
 ax.legend([h2], [i.get_label() for i in [h2]], loc="best")
-# -
-
+# %% [markdown]
 # Example with `min_area_threshold` set to 0.35, the intersection result in a cell should cover 35% or more of the cell area.
 
-# +
+# %%
 # min_area_threshold example
 
 result3 = ix.intersect(p, min_area_fraction=0.35)
@@ -212,21 +247,23 @@
     )
 # add legend
 ax.legend([h2], [i.get_label() for i in [h2]], loc="best")
-# -
-
+# %% [markdown]
 # ### <a id="rectgrid.2"></a>[Polyline with regular grid](#top)
 # MultiLineString to intersect with:
 
+# %%
 ls1 = LineString([(95, 105), (30, 50)])
 ls2 = LineString([(30, 50), (90, 22)])
 ls3 = LineString([(90, 22), (0, 0)])
 mls = MultiLineString(lines=[ls1, ls2, ls3])
 
+# %%
 result = ix.intersect(mls)
 
+# %% [markdown]
 # Plot the result
 
-# +
+# %%
 fig, ax = plt.subplots(1, 1, figsize=(8, 8))
 sgr.plot(ax=ax)
 ix.plot_intersection_result(result, ax=ax, cmap="viridis")
@@ -240,26 +277,28 @@
     )
 
 ax.legend([h2], [i.get_label() for i in [h2]], loc="best")
-# -
-
+# %% [markdown]
 # ### [MultiPoint with regular grid](#top)<a id="rectgrid.3"></a>
 #
 # MultiPoint to intersect with
 
+# %%
 mp = MultiPoint(
     points=[Point(50.0, 0.0), Point(45.0, 45.0), Point(10.0, 10.0), Point(150.0, 100.0)]
 )
 
+# %% [markdown]
 # For points and linestrings there is a keyword argument `return_all_intersections` which will return multiple intersection results for points or (parts of) linestrings on cell boundaries. As an example, the difference is shown with the MultiPoint intersection. Note the number of red "+" symbols indicating the centroids of intersected cells, in the bottom left case, there are 4 results because the point lies exactly on the intersection between 4 grid cells.
 
+# %%
 result = ix.intersect(mp)
 result_all = ix.intersect(mp, return_all_intersections=True)
 
-# +
+# %%
 fig, ax = plt.subplots(1, 1, figsize=(8, 8))
 sgr.plot(ax=ax)
-ix.plot_point(result, ax=ax, s=50, color="C0")
-ix.plot_point(result_all, ax=ax, s=50, marker=".", color="C3")
+ix.plot_point(result.ixshapes, ax=ax, ms=10, color="C0")
+ix.plot_point(result_all.ixshapes, ax=ax, ms=10, marker=".", color="C3")
 
 for irow, icol in result.cellids:
     (h2,) = ax.plot(
@@ -280,10 +319,54 @@
     )
 
 ax.legend([h2, h3], [i.get_label() for i in [h2, h3]], loc="best")
-# -
+# %% [markdown]
+# In addition to the `intersect()` and `intersects()` methods there is function to
+# quickly get the cellids for many points. This is done with `points_to_cellids()`.
+#
+# First lets create an array containing shapely Points.
+
+# %%
+n_pts = 10
+rng = np.random.default_rng(seed=42)
+x_coords = rng.uniform(0, 100, n_pts)
+y_coords = rng.uniform(0, 100, n_pts)
+random_points = shapely.points(x_coords, y_coords)
+
+# %% [markdown]
+# Now find the cellid for each point. `points_to_cellids` will only return a single result
+# for each point. In case a point is on the boundary between multiple cells, it will
+# return the cell with the lowest cellid.
+#
+#
+# In this example we're returning the result as a
+# node number to easily select the grid cells for our plot. But by default this method
+# returns (row, col) for structured grids.
 
+# %%
+# return cellid as node numbers: so (node,) instead of (row, col)
+# this makes it easier to select the grid cells to highlight in the plot
+result = ix.points_to_cellids(random_points, return_nodenumbers=True)
+
+# %% [markdown]
+# Plot the result
+
+# %%
+fig, ax = plt.subplots(1, 1, figsize=(8, 8))
+sgr.plot(ax=ax)
+ax.plot(x_coords, y_coords, "ro", ms=5, label="random points")
+ix.plot_polygon(
+    ix.geoms[result.cellids.astype(int)], ax=ax, fc="yellow", edgecolor="k", zorder=5
+)
+# %% [markdown]
+# Note that `points_to_cellids()` returns NaNs for points that lie outside the model grid.
+
+# %%
+ix.points_to_cellids(shapely.points([50, 110], [55, 50]), dataframe=True)
+
+# %% [markdown]
 # ## <a id="trigrid"></a>[Vertex Grid](#top)
 
+# %%
 cell2d = [
     [0, 83.33333333333333, 66.66666666666667, 3, 4, 2, 7],
     [1, 16.666666666666668, 33.333333333333336, 3, 4, 0, 5],
@@ -307,17 +390,20 @@
 ]
 tgr = fgrid.VertexGrid(vertices, cell2d)
 
+# %%
 fig, ax = plt.subplots(1, 1, figsize=(8, 8))
 pmv = fplot.PlotMapView(modelgrid=tgr)
 pmv.plot_grid(ax=ax)
-
+# %% [markdown]
 # ### <a id="trigrid.1"></a>[Polygon with triangular grid](#top)
 
+# %%
 ix2 = GridIntersect(tgr)
 
+# %%
 result = ix2.intersect(p)
 
-# +
+# %%
 fig, ax = plt.subplots(1, 1, figsize=(8, 8))
 ix.plot_intersection_result(result, ax=ax)
 
@@ -331,13 +417,13 @@
     )
 
 ax.legend([h2], [i.get_label() for i in [h2]], loc="best")
-# -
-
+# %% [markdown]
 # ### <a id="trigrid.2"></a>[LineString with triangular grid](#top)
 
+# %%
 result = ix2.intersect(mls)
 
-# +
+# %%
 fig, ax = plt.subplots(1, 1, figsize=(8, 8))
 ix2.plot_intersection_result(result, ax=ax, lw=3)
 
@@ -350,16 +436,16 @@
     )
 
 ax.legend([h2], [i.get_label() for i in [h2]], loc="best")
-# -
-
+# %% [markdown]
 # ### <a id="trigrid.3"></a>[MultiPoint with triangular grid](#top)
 
+# %%
 result = ix2.intersect(mp)
 result_all = ix2.intersect(mp, return_all_intersections=True)
 
-# +
+# %%
 fig, ax = plt.subplots(1, 1, figsize=(8, 8))
-ix2.plot_intersection_result(result, ax=ax, color="k", zorder=5, s=80)
+ix2.plot_intersection_result(result, ax=ax, color="k", zorder=5, ms=10)
 
 for cellid in result.cellids:
     (h2,) = ax.plot(
@@ -379,3 +465,5 @@
     )
 
 ax.legend([h2, h3], [i.get_label() for i in [h2, h3]], loc="best")
+
+# %%