diff --git a/autotest/t065_test_gridintersect.py b/autotest/t065_test_gridintersect.py index 3bb10540d2..c13c1bbf04 100644 --- a/autotest/t065_test_gridintersect.py +++ b/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: diff --git a/flopy/utils/gridintersect.py b/flopy/utils/gridintersect.py index 09137b9cf6..10a2c0f07f 100644 --- a/flopy/utils/gridintersect.py +++ b/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: