feat(GridIntersect): add optional kwarg to .intersect() to keep all c…

autotest/t065_test_gridintersect.py

@@ -86,12 +86,17 @@ def get_rect_grid(angrot=0.0, xyoffset=0.0, top=None, botm=None):
     return sgr
 
 
-def get_rect_vertex_grid(angrot=0.0, xyoffset=0.0):
+def get_rect_vertex_grid(angrot=0.0, xyoffset=0.0, top=None, botm=None):
+    # Creates a 2-layer, 2x2 rectangular vertex grid
     cell2d = [
         [0, 5.0, 5.0, 4, 0, 1, 4, 3],
         [1, 15.0, 5.0, 4, 1, 2, 5, 4],
         [2, 5.0, 15.0, 4, 3, 4, 7, 6],
         [3, 15.0, 15.0, 4, 4, 5, 8, 7],
+        [4, 5.0, 5.0, 4, 0, 1, 4, 3],
+        [5, 15.0, 5.0, 4, 1, 2, 5, 4],
+        [6, 5.0, 15.0, 4, 3, 4, 7, 6],
+        [7, 15.0, 15.0, 4, 4, 5, 8, 7],
     ]
     vertices = [
         [0, 0.0, 0.0],
@@ -107,8 +112,8 @@ def get_rect_vertex_grid(angrot=0.0, xyoffset=0.0):
     tgr = fgrid.VertexGrid(
         vertices,
         cell2d,
-        botm=np.atleast_2d(np.zeros(len(cell2d))),
-        top=np.ones(len(cell2d)),
+        top=top,
+        botm=botm,
         xoff=xyoffset,
         yoff=xyoffset,
         angrot=angrot,
@@ -317,7 +322,10 @@ def test_rect_vertex_grid_point_in_one_cell_shapely(rtree=True):
         import shapely
     except:
         return
-    gr = get_rect_vertex_grid()
+    botm = np.concatenate([np.ones(4), 0.5 * np.ones(4), np.zeros(4)]).reshape(
+        3, 2, 2
+    )
+    gr = get_rect_vertex_grid(top=np.ones(4), botm=botm)
     ix = GridIntersect(gr, method="vertex", rtree=rtree)
     result = ix.intersect(Point(4.0, 4.0))
     assert len(result) == 1
@@ -334,6 +342,25 @@ def test_rect_vertex_grid_point_in_one_cell_shapely(rtree=True):
     return result
 
 
+def test_rect_vertex_grid_point_in_one_cell_shapely_all_layers(rtree=True):
+    # avoid test fail when shapely not available
+    try:
+        import shapely
+    except:
+        return
+    # botm = np.concatenate([np.ones(4), np.zeros(4)]).reshape(2, 2, 2)
+    botm = np.concatenate([np.ones(4), 0.5 * np.ones(4), np.zeros(4)]).reshape(
+        3, 2, 2
+    )
+    gr = get_rect_vertex_grid(top=np.ones(4), botm=botm)
+    ix = GridIntersect(gr, method="vertex", rtree=rtree)
+    result = ix.intersect(Point(4.0, 4.0), keep_stacked_cells=True)
+    assert len(result) == 2
+    assert result.cellids[0] == 0
+    assert result.cellids[1] == 4
+    return result
+
+
 def test_rect_grid_multipoint_in_one_cell_shapely(rtree=True):
     # avoid test fail when shapely not available
     try:
@@ -595,6 +622,26 @@ def test_rect_grid_linestring_in_2cells_shapely(rtree=True):
     return result
 
 
+def test_rect_grid_linestring_in_2cells_shapely_all_layers(rtree=True):
+    # avoid test fail when shapely not available
+    try:
+        import shapely
+    except:
+        return
+    botm = np.concatenate([np.ones(4), 0.5 * np.ones(4), np.zeros(4)]).reshape(
+        3, 2, 2
+    )
+    gr = get_rect_vertex_grid(top=np.ones(4), botm=botm)
+    ix = GridIntersect(gr, method="vertex", rtree=rtree)
+    result = ix.intersect(LineString([(5.0, 5.0), (15.0, 5.0)]), keep_stacked_cells=True)
+    assert len(result) == 4
+    assert result.lengths.sum() == 20.0
+    assert result.cellids[0] == 0
+    assert result.cellids[1] == 1
+    assert result.cellids[2] == 4
+    assert result.cellids[3] == 5
+
+
 def test_rect_grid_linestring_on_outer_boundary_shapely(rtree=True):
     # avoid test fail when shapely not available
     try:

flopy/utils/gridintersect.py

@@ -223,6 +223,8 @@ def intersect(self, shp, **kwargs):
         keepzerolengths : bool
             boolean method to keep zero length intersections for
             linestring intersection
+        keep_stacked_cells : bool
+            For shapely only, keeps all intersected cells even if cell vertices are already in the results.
 
         Returns
         -------
@@ -233,6 +235,7 @@ def intersect(self, shp, **kwargs):
         shp = gu.shapely
         sort_by_cellid = kwargs.pop("sort_by_cellid", True)
         keepzerolengths = kwargs.pop("keepzerolengths", False)
+        keep_stacked_cells = kwargs.pop("keep_stacked_cells", False)
 
         if gu.shapetype in ("Point", "MultiPoint"):
             if (
@@ -241,7 +244,9 @@ def intersect(self, shp, **kwargs):
             ):
                 rec = self._intersect_point_structured(shp)
             else:
-                rec = self._intersect_point_shapely(shp, sort_by_cellid)
+                rec = self._intersect_point_shapely(
+                    shp, sort_by_cellid, keep_stacked_cells
+                )
         elif gu.shapetype in ("LineString", "MultiLineString"):
             if (
                 self.method == "structured"
@@ -252,7 +257,7 @@ def intersect(self, shp, **kwargs):
                 )
             else:
                 rec = self._intersect_linestring_shapely(
-                    shp, keepzerolengths, sort_by_cellid
+                    shp, keepzerolengths, sort_by_cellid, keep_stacked_cells
                 )
         elif gu.shapetype in ("Polygon", "MultiPolygon"):
             if (
@@ -456,7 +461,9 @@ def sort_key(o):
         shapelist.sort(key=sort_key)
         return shapelist
 
-    def _intersect_point_shapely(self, shp, sort_by_cellid=True):
+    def _intersect_point_shapely(
+        self, shp, sort_by_cellid=True, keep_stacked_cells=False
+    ):
         """intersect grid with Point or MultiPoint.
 
         Parameters
@@ -468,6 +475,8 @@ def _intersect_point_shapely(self, shp, sort_by_cellid=True):
         sort_by_cellid : bool, optional
             flag whether to sort cells by id, used to ensure node
             with lowest id is returned, by default True
+        keep_stacked_cells : bool, optional
+            For shapely only, keeps all intersected cells even if cell vertices are already in the results.
 
         Returns
         -------
@@ -505,7 +514,7 @@ def _intersect_point_shapely(self, shp, sort_by_cellid=True):
             for c in collection:
                 verts = c.__geo_interface__["coordinates"]
                 # avoid returning multiple cells for points on boundaries
-                if verts in parsed_points:
+                if not keep_stacked_cells and verts in parsed_points:
                     continue
                 parsed_points.append(verts)
                 cell_shps.append(c)  # collect only new points
@@ -534,7 +543,11 @@ def _intersect_point_shapely(self, shp, sort_by_cellid=True):
         return rec
 
     def _intersect_linestring_shapely(
-        self, shp, keepzerolengths=False, sort_by_cellid=True
+        self,
+        shp,
+        keepzerolengths=False,
+        sort_by_cellid=True,
+        keep_stacked_cells=False,
     ):
         """intersect with LineString or MultiLineString.
 
@@ -547,6 +560,8 @@ def _intersect_linestring_shapely(
         sort_by_cellid : bool, optional
             flag whether to sort cells by id, used to ensure node
             with lowest id is returned, by default True
+        keep_stacked_cells : bool, optional
+            For shapely only, keeps all intersected cells even if cell vertices are already in the results.
 
         Returns
         -------
@@ -580,7 +595,7 @@ def _intersect_linestring_shapely(
             for c in collection:
                 verts = c.__geo_interface__["coordinates"]
                 # test if linestring was already processed (if on boundary)
-                if verts in vertices:
+                if not keep_stacked_cells and verts in vertices:
                     continue
                 # if keep zero don't check length
                 if not keepzerolengths: