From df568ace01d2dcfdb86da23b423b28e03e74fa69 Mon Sep 17 00:00:00 2001
From: Huite Bootsma <huitebootsma@gmail.com>
Date: Thu, 3 Aug 2023 18:28:57 +0200
Subject: [PATCH] Add snap_to_grid to user guide. Fix shapely_index versus
 line_index.

---
 examples/vector_conversion.py | 28 ++++++++++++++++++++++++++++
 xugrid/ugrid/snapping.py      | 15 ++++++++++-----
 2 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/examples/vector_conversion.py b/examples/vector_conversion.py
index 7a8cc6d50..e5ca75dcb 100644
--- a/examples/vector_conversion.py
+++ b/examples/vector_conversion.py
@@ -207,3 +207,31 @@
 #   result will essentially be one polygon per face. In such cases, it is more
 #   efficient to use ``xugrid.UgridDataArray.to_geodataframe``, which directly
 #   converts every face to a polygon.
+#
+# Snap to grid
+# ------------
+#
+# The examples above deal with data on the faces of the grid. However, data can
+# also be associated with the nodes or edges of the grid. For example, we can
+# snap line data to exactly follow the edges (the face boundaries):
+
+linestrings = gpd.GeoDataFrame(geometry=utrecht.exterior)
+snapped_uds, snapped_gdf = xu.snap_to_grid(linestrings, uda, max_snap_distance=1.0)
+
+fig, ax = plt.subplots()
+snapped_gdf.plot(ax=ax)
+snapped_uds.ugrid.grid.plot(ax=ax, edgecolor="gray", linewidth=0.5)
+utrecht.plot(ax=ax, edgecolor="red", facecolor="none", linewidth=0.75)
+ax.set_xlim(xmin, xmax)
+ax.set_ylim(ymin, ymax)
+
+# %%
+# There are (arguably) many ways in which a linestring could be snapped to
+# edges. The function above uses the criterion the following criterion: if a
+# part of the linestring is located between the centroid of a face and the
+# centroid of an edge, it is snapped to that edge.
+#
+# This sometimes result in less (visually) pleasing results, such as the two
+# triangles in the lower left corner which are surrounded by the snapped edges.
+# In general, results are best when the line segments are relatively large
+# compare to the grid faces.
diff --git a/xugrid/ugrid/snapping.py b/xugrid/ugrid/snapping.py
index 242aa9f92..2040572fc 100644
--- a/xugrid/ugrid/snapping.py
+++ b/xugrid/ugrid/snapping.py
@@ -334,12 +334,12 @@ def _create_output_gdf(
     vertices,
     edge_node_connectivity,
     edges,
-    line_index,
+    shapely_index,
 ):
     edge_vertices = vertices[edge_node_connectivity[edges]]
     geometry = shapely.linestrings(edge_vertices)
     return gpd.GeoDataFrame(
-        lines.drop(columns="geometry").iloc[line_index], geometry=geometry
+        lines.drop(columns="geometry").iloc[shapely_index], geometry=geometry
     )
 
 
@@ -401,12 +401,14 @@ def snap_to_grid(
     # Create geometric data
     edge_centroids = vertices[edge_node_connectivity].mean(axis=1)
     line_geometry = coerce_geometry(lines)
-    line_coords, line_index = shapely.get_coordinates(line_geometry, return_index=True)
+    line_coords, shapely_index = shapely.get_coordinates(
+        line_geometry, return_index=True
+    )
     # Snap line_coords to grid
     x, y = snap_to_nodes(
         *line_coords.T, *vertices.T, max_snap_distance, tiebreaker="nearest"
     )
-    line_edges = lines_as_edges(np.column_stack([x, y]), line_index)
+    line_edges = lines_as_edges(np.column_stack([x, y]), shapely_index)
 
     # Search for intersections. Every edge is potentially divided into smaller
     # segments: The segment_indices contain (repeated) values of the
@@ -443,7 +445,10 @@ def snap_to_grid(
     # When multiple line parts are snapped to the same edge, use the ones with
     # the greatest length inside the cell.
     edges, line_index = _find_largest_edges(segment_edges, edge_index, line_index)
+    shapely_index = shapely_index[line_index]
 
     uds = _create_output_dataset(lines, topology, edges, line_index)
-    gdf = _create_output_gdf(lines, vertices, edge_node_connectivity, edges, line_index)
+    gdf = _create_output_gdf(
+        lines, vertices, edge_node_connectivity, edges, shapely_index
+    )
     return uds, gdf