Repository: Deltares/hydromt

Issue:
ENH: improve speed of raster.vector_grid method.
### Feature Type

- Changing existing functionality in HydroMT

### Enhancement Description

The current vector grid implementation is slow. This is used to visualize the grid as well as for some GIS methods, hence speed is important.

### Feature Description

This can be enhanced by avoiding looping over each cell. In addition, instead of a polygon per cell only the lines of the cell edges could be returned. 

### Additional Context

_No response_

Discussion:

Patch:
diff --git a/docs/changelog.rst b/docs/changelog.rst
index 9dcad562..def2d567 100644
--- a/docs/changelog.rst
+++ b/docs/changelog.rst
@@ -14,6 +14,7 @@ Added
 
 Changed
 -------
+- new geom_type argument in `RasterDataArray.vector_grid` to specify the geometry type {'polygon', 'line', 'points'} of the vector grid. PR #351
 - Added extrapolate option to `raster.interpolate_na` method. PR #348
 
 Fixed
diff --git a/hydromt/raster.py b/hydromt/raster.py
index f641ddf7..78d471fd 100644
--- a/hydromt/raster.py
+++ b/hydromt/raster.py
@@ -36,7 +36,8 @@ from rasterio.enums import Resampling
 from scipy import ndimage
 from scipy.interpolate import griddata
 from scipy.spatial import cKDTree
-from shapely.geometry import Polygon, box
+from shapely.geometry import Polygon, box, LineString
+import shapely
 
 from . import _compat, gis_utils
 
@@ -1214,18 +1215,57 @@ class XRasterBase(XGeoBase):
         da_out.attrs.pop("_FillValue", None)
         return da_out.astype(bool)
 
-    def vector_grid(self):
-        """Return a geopandas GeoDataFrame with a geometry for each grid cell."""
-        transform = self.transform
+    def vector_grid(self, geom_type: str = "polygon") -> gpd.GeoDataFrame:
+        """Return a vector representation of the grid.
+
+        Parameters
+        ----------
+        geom_type : str, optional
+            Type of geometry to return, by default 'polygon'
+            Available options are 'polygon', 'line', 'point'
+
+        Returns
+        -------
+        geopandas.GeoDataFrame
+            GeoDataFrame with grid geometries
+        """
+        if not hasattr(shapely, "from_ragged_array"):
+            raise ImportError("the vector_grid method requires shapely 2.0+")
         nrow, ncol = self.shape
-        cells = []
-        for i in range(nrow):
-            rs = np.array([i, i + 1, i + 1, i, i])
-            for j in range(ncol):
-                cs = np.array([j, j, j + 1, j + 1, j])
-                xs, ys = transform * (cs, rs)
-                cells.append(Polygon([*zip(xs, ys)]))
-        return gpd.GeoDataFrame(geometry=cells, crs=self.crs)
+        transform = self.transform
+        if geom_type.lower().startswith("polygon"):
+            # build a flattened list of coordinates for each cell
+            dr = np.array([0, 0, 1, 1, 0])
+            dc = np.array([0, 1, 1, 0, 0])
+            ii, jj = np.meshgrid(np.arange(0, nrow), np.arange(0, ncol))
+            coords = np.empty((nrow * ncol * 5, 2), dtype=np.float64)
+            # order of cells: first rows, then cols
+            coords[:, 0], coords[:, 1] = transform * (
+                (jj.T[:, :, None] + dr[None, None, :]).ravel(),
+                (ii.T[:, :, None] + dc[None, None, :]).ravel(),
+            )
+            # offsets of the first coordinate of each cell, index of cells
+            offsets = (
+                np.arange(0, nrow * ncol * 5 + 1, 5, dtype=np.int32),
+                np.arange(nrow * ncol + 1, dtype=np.int32),
+            )
+            # build the geometry array in a fast way
+            geoms = shapely.from_ragged_array(3, coords, offsets)  # type 3 is polygon
+        elif geom_type.lower().startswith("line"):  # line or linestring
+            geoms = []
+            # horizontal lines
+            for i in range(nrow + 1):
+                xs, ys = transform * (np.array([0, ncol]), np.array([i, i]))
+                geoms.append(LineString(zip(xs, ys)))
+            for i in range(ncol + 1):
+                xs, ys = transform * (np.array([i, i]), np.array([0, nrow]))
+                geoms.append(LineString(zip(xs, ys)))
+        elif geom_type.lower().startswith("point"):
+            x, y = gis_utils.affine_to_meshgrid(transform, (nrow, ncol))
+            geoms = gpd.points_from_xy(x.ravel(), y.ravel())
+        else:
+            raise ValueError(f"geom_type {geom_type} not recognized")
+        return gpd.GeoDataFrame(geometry=geoms, crs=self.crs)
 
     def area_grid(self, dtype=np.float32):
         """Returns the grid cell area [m2].


Test patch:
diff --git a/tests/test_raster.py b/tests/test_raster.py
index c2ff47ac..46f57afc 100644
--- a/tests/test_raster.py
+++ b/tests/test_raster.py
@@ -354,9 +354,22 @@ def test_interpolate_na():
 
 
 def test_vector_grid(rioda):
+    # polygon
     gdf = rioda.raster.vector_grid()
-    assert isinstance(gdf, gpd.GeoDataFrame)
     assert rioda.raster.size == gdf.index.size
+    assert np.all(gdf.geometry.type == "Polygon")
+    assert np.all(gdf.total_bounds == rioda.raster.bounds)
+    # line
+    gdf = rioda.raster.vector_grid(geom_type="line")
+    nrow, ncol = rioda.raster.shape
+    assert gdf.index.size == (nrow + 1 + ncol + 1)
+    assert np.all(gdf.geometry.type == "LineString")
+    assert np.all(gdf.total_bounds == rioda.raster.bounds)
+    # point
+    gdf = rioda.raster.vector_grid(geom_type="point")
+    assert np.all(gdf.geometry.type == "Point")
+    assert rioda.raster.size == gdf.index.size
+    assert np.all(gdf.intersects(rioda.raster.box.geometry[0]))
 
 
 def test_sample():