Merge pull request #283 from Deltares/improve-snapping

Improve snapping

docs/api.rst

@@ -142,6 +142,7 @@ nodes/edges/faces if within a certain range.
 .. autosummary::
       :toctree: api/
 
+       snap_nodes
        snap_to_grid
        create_snap_to_grid_dataframe
 

docs/changelog.rst

@@ -22,6 +22,17 @@ Fixed
   rather than a single delta to denote cell sizes along a dimension (i.e.
   ``dy`` along ``y`` midpoints, and ``dx`` along ``x``.)
 
+Added
+~~~~~
+
+- :func:`xugrid.snap_nodes` to snap neigbhoring vertices together that are
+  located within a maximum snapping distance from each other. If vertices are
+  located within a maximum distance, some of them are snapped to their
+  neighbors ("targets"), thereby guaranteeing a minimum distance between nodes
+  in the result. The determination of whether a point becomes a target itself
+  or gets snapped to another point is primarily based on the order in which
+  points are processed and their spatial relationships.
+
 [0.11.1] 2024-08-13
 -------------------
 

tests/test_snap.py

@@ -34,7 +34,32 @@ def structured():
     return xr.DataArray(np.ones(shape, dtype=np.int32), coords=coords, dims=dims)
 
 
-def test_snap__three_points():
+def test_snap__three_points_horizontal():
+    x = np.array([0.0, 1.0, 2.0])
+    y = np.zeros_like(x)
+    inv_perm, snap_x, snap_y = snap_nodes(x, y, 0.1)
+    assert inv_perm is None
+    assert np.array_equal(x, snap_x)
+    assert np.array_equal(y, snap_y)
+
+    inv_perm, snap_x, snap_y = snap_nodes(x, y, 1.0)
+    expected_inv_perm = np.array([0, 0, 1])
+    expected_x = np.array([0.0, 2.0])
+    expected_y = np.zeros_like(expected_x)
+    assert np.array_equal(inv_perm, expected_inv_perm)
+    assert np.array_equal(expected_x, snap_x)
+    assert np.array_equal(expected_y, snap_y)
+
+    inv_perm, snap_x, snap_y = snap_nodes(x, y, 2.0)
+    expected_inv_perm = np.array([0, 0, 0])
+    expected_x = np.array([0.0])
+    expected_y = np.zeros_like(expected_x)
+    assert np.array_equal(inv_perm, expected_inv_perm)
+    assert np.array_equal(expected_x, snap_x)
+    assert np.array_equal(expected_y, snap_y)
+
+
+def test_snap__three_points_diagonal():
     x = y = np.array([0.0, 1.0, 1.5])
     inv_perm, snap_x, snap_y = snap_nodes(x, y, 0.1)
     assert inv_perm is None
@@ -44,17 +69,18 @@ def test_snap__three_points():
     # hypot(0.5, 0.5) = 0.707...
     inv_perm, snap_x, snap_y = snap_nodes(x, y, 0.71)
     expected_inv_perm = np.array([0, 1, 1])
-    expected_x = expected_y = np.array([0.0, 1.25])
+    expected_x = expected_y = np.array([0.0, 1.0])
     assert np.array_equal(inv_perm, expected_inv_perm)
     assert np.array_equal(snap_x, expected_x)
     assert np.array_equal(snap_y, expected_y)
 
     # hypot(1, 1) = 1.414...
     inv_perm, snap_x, snap_y = snap_nodes(x, y, 1.42)
-    expected_inv_perm = np.array([0, 0, 0])
+    expected_inv_perm = np.array([0, 1, 1])
+    expected_x = expected_y = np.array([0.0, 1.5])
     assert np.array_equal(inv_perm, expected_inv_perm)
-    assert np.allclose(snap_x, np.array([2.5 / 3]))
-    assert np.allclose(snap_y, np.array([2.5 / 3]))
+    assert np.array_equal(snap_x, expected_x)
+    assert np.array_equal(snap_y, expected_y)
 
 
 def test_snap__two_lines():
@@ -70,7 +96,7 @@ def test_snap__two_lines():
     c = inv_perm[edge_node_connectivity]
 
     expected_inv_perm = np.array([0, 1, 1, 2])
-    expected_x = np.array([0.0, 1.01, 2.0])
+    expected_x = np.array([0.0, 1.0, 2.0])
     expected_y = np.array([1.0, 0.0, 1.0])
     expected_c = np.array(
         [

tests/test_sparse.py

@@ -1,9 +1,15 @@
+import os
+
 import numpy as np
 import pytest
 
 from xugrid.core import sparse
 
 
+def numba_enabled() -> bool:
+    return os.environ.get("NUMBA_DISABLE_JIT") != "1"
+
+
 @pytest.fixture(scope="function")
 def coo_matrix():
     source_index = np.arange(10)
@@ -37,13 +43,41 @@ def test_weight_matrix_csr(csr_matrix):
     assert csr_matrix.nnz == 10
 
 
+@pytest.mark.skipif(
+    numba_enabled(),
+    reason="Numba cannot convert native range_state_int64 to Python object.",
+)
 def test_nzrange(csr_matrix):
+    # These functions work fine if called inside of other numba functions when
+    # numba is enabled.
     i = sparse.nzrange(csr_matrix, 0)
     assert np.array_equal(i, range(0, 2))
     i = sparse.nzrange(csr_matrix, 1)
     assert np.array_equal(i, range(2, 4))
 
 
+@pytest.mark.skipif(
+    numba_enabled(),
+    reason="Function returns a slice object; python and no-python slices don't mix.",
+)
+def test_row_slice(csr_matrix):
+    # These functions work fine if called inside of other numba functions when
+    # numba is enabled.
+    assert sparse.row_slice(csr_matrix, 0) == slice(0, 2, None)
+
+
+@pytest.mark.skipif(
+    numba_enabled(),
+    reason="Function returns a zip object; python and no-python zips don't mix.",
+)
+def test_columns_and_values(csr_matrix):
+    # These functions work fine if called inside of other numba functions when
+    # numba is enabled.
+    zipped = sparse.columns_and_values(csr_matrix, sparse.row_slice(csr_matrix, 0))
+    result = list(zipped)
+    assert result == [(0, 0.5), (1, 0.5)]
+
+
 def test_coo_to_csr(coo_matrix):
     csr_matrix = coo_matrix.to_csr()
     assert isinstance(csr_matrix, sparse.MatrixCSR)

xugrid/__init__.py

@@ -24,7 +24,11 @@
 from xugrid.ugrid.conventions import UgridRolesAccessor
 from xugrid.ugrid.partitioning import merge_partitions
 from xugrid.ugrid.polygonize import polygonize
-from xugrid.ugrid.snapping import create_snap_to_grid_dataframe, snap_to_grid
+from xugrid.ugrid.snapping import (
+    create_snap_to_grid_dataframe,
+    snap_nodes,
+    snap_to_grid,
+)
 from xugrid.ugrid.ugrid1d import Ugrid1d
 from xugrid.ugrid.ugrid2d import Ugrid2d
 
@@ -55,6 +59,7 @@
     "UgridRolesAccessor",
     "merge_partitions",
     "polygonize",
+    "snap_nodes",
     "snap_to_grid",
     "create_snap_to_grid_dataframe",
     "Ugrid1d",

xugrid/core/sparse.py

@@ -146,3 +146,8 @@ def row_slice(A, row: int) -> slice:
     start = A.indptr[row]
     end = A.indptr[row + 1]
     return slice(start, end)
+
+
+@numba.njit(inline="always")
+def columns_and_values(A, slice):
+    return zip(A.indices[slice], A.data[slice])

xugrid/ugrid/connectivity.py

@@ -255,7 +255,7 @@ def ragged_index(n: int, m: int, m_per_row: IntArray) -> BoolArray:
     m = 4
     m_per_row = np.array([1, 2, 3])
 
-    Then the result of _ragged_index(n, m, m_per_row) is:
+    Then the result of ragged_index(n, m, m_per_row) is:
 
     np.array([
         [True, False, False, False],

xugrid/ugrid/interpolate.py

@@ -8,7 +8,7 @@
 from scipy import sparse
 
 from xugrid.constants import FloatArray, IntArray
-from xugrid.core.sparse import MatrixCSR, nzrange, row_slice
+from xugrid.core.sparse import MatrixCSR, columns_and_values, nzrange, row_slice
 
 
 @nb.njit(inline="always")
@@ -21,11 +21,6 @@ def upper_slice(ilu, row: int) -> slice:
     return slice(ilu.uptr[row], ilu.indptr[row + 1])
 
 
-@nb.njit(inline="always")
-def columns_and_values(ilu, slice):
-    return zip(ilu.indices[slice], ilu.data[slice])
-
-
 @nb.njit
 def set_uptr(ilu: ILU0Preconditioner) -> None:
     # i is row index, j is column index

xugrid/ugrid/snapping.py

@@ -5,8 +5,6 @@
 import numpy as np
 import pandas as pd
 import xarray as xr
-from scipy import sparse
-from scipy.sparse.csgraph import connected_components
 from scipy.spatial import cKDTree
 
 import xugrid as xu
@@ -19,6 +17,7 @@
     Point,
     Vector,
 )
+from xugrid.core.sparse import MatrixCSR, columns_and_values, row_slice
 from xugrid.ugrid import connectivity
 from xugrid.ugrid.connectivity import AdjacencyMatrix
 from xugrid.ugrid.ugrid2d import Ugrid2d
@@ -38,20 +37,58 @@
     shapely = MissingOptionalModule("shapely")
 
 
+@nb.njit(cache=True)
+def _snap_to_nearest(A: MatrixCSR, snap_candidates: IntArray, max_distance) -> IntArray:
+    """
+    Find a closest target for each node.
+
+    The kD tree distance matrix will have stored for each node the other nodes
+    that are within snapping distance. These are the rows in the sparse matrix
+    that have more than one entry: the snap_candidates.
+
+    The first condition for a point to become a TARGET is if it hasn't been
+    connected to another point yet, i.e. it is UNVISITED. Once a point becomes
+    an TARGET, it looks for nearby points within the max_distance. These nearby
+    points are connected if: they are UNVISITED (i.e. they don't have a target
+    yet), or the current target is closer than the previous.
+    """
+    UNVISITED = -1
+    TARGET = -2
+    nearest = np.full(A.n, max_distance + 1.0)
+    visited = np.full(A.n, UNVISITED)
+
+    for i in snap_candidates:
+        if visited[i] != UNVISITED:
+            continue
+        visited[i] = TARGET
+
+        # Now iterate through every node j that is within max_distance of node i.
+        for j, dist in columns_and_values(A, row_slice(A, i)):
+            if i == j or visited[j] == TARGET:
+                # Continue if we're looking at the distance to ourselves
+                # (i==j), or other node is a target.
+                continue
+            if visited[j] == UNVISITED or dist < nearest[j]:
+                # If unvisited node, or already visited but we're closer, set
+                # to i.
+                visited[j] = i
+                nearest[j] = dist
+
+    return visited
+
+
 def snap_nodes(
-    x: FloatArray, y: FloatArray, max_distance: float
+    x: FloatArray, y: FloatArray, max_snap_distance: float
 ) -> Tuple[FloatArray, FloatArray, IntArray]:
     """
     Snap neigbhoring vertices together that are located within a maximum
-    distance from each other.
-
-    If vertices are located within a maximum distance, they are merged into a
-    single vertex. The coordinates of the merged coordinates are given by the
-    centroid of the merging vertices.
+    snapping distance from each other.
 
-    Note that merging is a communicative process: vertex A might lie close to
-    vertex B, vertex B might lie close to vertex C, and so on. These points are
-    grouped and merged into a single new vertex.
+    If vertices are located within a maximum distance, some of them are snapped
+    to their neighbors ("targets"), thereby guaranteeing a minimum distance
+    between nodes in the result. The determination of whether a point becomes a
+    target itself or gets snapped to another point is primarily based on the
+    order in which points are processed and their spatial relationships.
 
     This function also return an inverse index array. In case of a connectivity
     array, ``inverse`` can be used to index into, yielding the updated
@@ -63,54 +100,39 @@ def snap_nodes(
     ----------
     x: 1D nd array of floats of size N
     y: 1D nd array of floats of size N
-    max_distance: float
+    max_snap_distance: float
 
     Returns
     -------
     inverse: 1D nd array of ints of size N
         Inverse index array: the new vertex number for every old vertex. Is
         None when no vertices within max_distance of each other.
-    x_merged: 1D nd array of floats of size M
+    x_snapped: 1D nd array of floats of size M
         Returns a copy of ``x`` when no vertices within max_distance of each
         other.
-    y_merged: 1D nd array of floats of size M
+    y_snapped: 1D nd array of floats of size M
         Returns a copy of ``y`` when no vertices within max_distance of each
         other.
     """
     # First, find all the points that lie within max_distance of each other
     coords = np.column_stack((x, y))
-    n = len(coords)
     tree = cKDTree(coords)
-    coo_distances = tree.sparse_distance_matrix(
-        tree, max_distance=max_distance, output_type="coo_matrix"
-    )
-    # Get rid of diagonal
-    i = coo_distances.row
-    j = coo_distances.col
-    off_diagonal = i != j
-    i = i[off_diagonal]
-    j = j[off_diagonal]
-    # If no pairs are found, there is nothing to snap
-    if len(i) > 0:
-        # Next, group the points together.
-        # Point A might lie close to point B, point B might lie close to
-        # Point C, and so on. These points are grouped, and their centroid
-        # is computed.
-        coo_content = (np.ones(i.size), (i, j))
-        coo_matrix = sparse.coo_matrix(coo_content, shape=(n, n))
-        # Directed is true: this matrix is symmetrical
-        _, inverse = connected_components(coo_matrix, directed=True)
-        new = (
-            pd.DataFrame({"label": inverse, "x": x, "y": y})
-            .groupby("label")
-            .agg(
-                {
-                    "x": "mean",
-                    "y": "mean",
-                }
-            )
+    distances = tree.sparse_distance_matrix(
+        tree, max_distance=max_snap_distance, output_type="coo_matrix"
+    ).tocsr()
+    should_snap = distances.getnnz(axis=1) > 1
+
+    if should_snap.any():
+        index = np.arange(x.size)
+        visited = _snap_to_nearest(
+            A=MatrixCSR.from_csr_matrix(distances),
+            snap_candidates=index[should_snap],
+            max_distance=max_snap_distance,
         )
-        return inverse, new["x"].to_numpy(), new["y"].to_numpy()
+        targets = visited < 0  # i.e. still UNVISITED or TARGET valued.
+        visited[targets] = index[targets]
+        deduplicated, inverse = np.unique(visited, return_inverse=True)
+        return inverse, x[deduplicated], y[deduplicated]
     else:
         return None, x.copy(), y.copy()