Reindex

docs/changelog.rst

@@ -23,6 +23,11 @@ Added
   ordinary grid to a periodic grid.
 - :attr:`xugrid.Ugrid2d.perimeter` has been added the compute the length of the
   face perimeters.
+- :meth:`xugrid.Ugrid1d.reindex_like`,
+  :meth:`xugrid.Ugrid2d.reindex_like`,
+  :meth:`xugrid.UgridDataArrayAccessor.reindex_like` and
+  :meth:`xugrid.UgridDatasetAccessor.reindex_like` have been added to deal with
+  equivalent but differently ordered topologies and data.
 
 Changed
 ~~~~~~~
@@ -33,6 +38,14 @@ Changed
 Fixed
 ~~~~~
 
+- Using an index which only reorders but does not change the size in
+  :meth:`xugrid.Ugrid1d.topology_subset` or
+  :meth:`xugrid.Ugrid2d.topology_subset` would erroneously result in the
+  original grid being returned, rather than a new grid with the faces or edges
+  shuffled. This breaks the link the between topology and data when using
+  ``.isel`` on a UgridDataset or UgridDataArray. This has been fixed: both data
+  and the topology are now shuffled accordingly. 
+
 [0.6.5] 2023-09-30
 ------------------
 

examples/partitioning.py

@@ -110,8 +110,18 @@
 
 # %%
 # The topology is equivalent, but the nodes, edges, and faces are in a
-# different order. This is because ``merge_partitions`` concatenates the
-# partitions. To preserve order, we can assign an ID to the partitions, and
+# different order. This is because ``merge_partitions`` simply concatenates the
+# partitions.
+#
+# The easiest way to restore the order is by providing an example of the
+# original topology. ``reindex_like`` looks at the coordinates of both
+# (equivalent!) grids and automatically determines how to reorder:
+
+reordered = merged.ugrid.reindex_like(uda)
+uda == reordered
+
+# %%
+# Alternatively, we can also assign IDs, carry these along, and use these to
 # reorder the data after merging.
 
 uds = xu.UgridDataset(grids=[uda.ugrid.grid])

tests/test_connectivity.py

@@ -31,6 +31,44 @@ def mixed_mesh():
     return faces, fill_value
 
 
+def test_argsort_rows():
+    with pytest.raises(ValueError, match="Array is not 2D"):
+        connectivity.argsort_rows(np.array([3, 2, 1, 0]))
+
+    array = np.array(
+        [
+            [1, 0],
+            [0, 1],
+            [2, 2],
+            [2, 1],
+            [0, 2],
+            [2, 0],
+        ]
+    )
+    _, expected = np.unique(array, axis=0, return_index=True)
+    actual = connectivity.argsort_rows(array)
+    assert np.array_equal(actual, expected)
+
+
+def test_index_like():
+    xy_a = np.array([[0.0, 0.0], [1.0, 1.0]])
+    xy_b = np.array([[0.0, 0.0]])
+    with pytest.raises(ValueError, match="coordinates do not match in shape"):
+        connectivity.index_like(xy_a, xy_b, tolerance=0.0)
+
+    xy_b = np.array([[0.0, 0.0], [1.1, 1.0]])
+    with pytest.raises(ValueError, match="coordinates are not identical after sorting"):
+        connectivity.index_like(xy_a, xy_b, tolerance=0.0)
+    # Now with higher tolerance
+    connectivity.index_like(xy_a, xy_b, tolerance=0.2)
+
+    xy_a = np.array([[3.0, 3.0], [1.0, 1.0], [2.0, 2.0], [0.0, 0.0]])
+    xy_b = np.array([[0.0, 0.0], [1.0, 1.0], [3.0, 3.0], [2.0, 2.0]])
+    actual = connectivity.index_like(xy_a, xy_b, tolerance=0.0)
+    expected = [3, 1, 0, 2]
+    assert np.array_equal(actual, expected)
+
+
 def test_neighbors():
     i = [0, 0, 0, 1, 1, 1]
     j = [0, 1, 2, 1, 3, 2]

tests/test_ugrid1d.py

@@ -330,6 +330,15 @@ def test_topology_subset():
     assert np.array_equal(actual.node_y, [1.0, 2.0])
 
 
+def test_reindex_like():
+    grid = grid1d()
+    index = np.array([1, 0])
+    reordered = grid.topology_subset(index)
+    obj = xr.DataArray(index, dims=(reordered.edge_dimension))
+    reindexed = reordered.reindex_like(grid, obj=obj)
+    assert np.array_equal(reindexed, [0, 1])
+
+
 def test_ugrid1d_plot():
     grid = grid1d()
     primitive = grid.plot()

tests/test_ugrid2d.py

@@ -846,6 +846,14 @@ def test_topology_subset():
     actual = grid.topology_subset(face_index)
     assert actual is grid
 
+    # Reordering
+    face_index = np.array([3, 2, 1, 0])
+    actual = grid.topology_subset(face_index)
+    assert np.array_equal(
+        actual.face_node_connectivity,
+        grid.face_node_connectivity[::-1],
+    )
+
     # Check that alternative attrs are preserved.
     grid = grid2d(
         attrs={"node_dimension": "nNetNode"},
@@ -856,6 +864,26 @@ def test_topology_subset():
     assert actual.node_dimension == "nNetNode"
 
 
+def test_reindex_like():
+    grid = grid2d()
+    # Change face and edge_index.
+    index = np.arange(grid.n_face)
+    rev_index = index[::-1]
+    edge_index = np.arange(grid.n_edge)
+    rev_edge_index = edge_index[::-1]
+    reordered = grid.topology_subset(rev_index)
+    reordered._edge_node_connectivity = reordered._edge_node_connectivity[
+        rev_edge_index
+    ]
+
+    face_da = xr.DataArray(rev_index, dims=(reordered.face_dimension,))
+    edge_da = xr.DataArray(rev_edge_index, dims=(reordered.edge_dimension,))
+    ds = xr.Dataset({"edge": edge_da, "face": face_da})
+    reindexed = reordered.reindex_like(grid, obj=ds)
+    assert np.array_equal(reindexed["face"], index)
+    assert np.array_equal(reindexed["edge"], edge_index)
+
+
 def test_triangulate():
     grid = grid2d()
     actual = grid.triangulate()

tests/test_ugrid_dataset.py

@@ -293,6 +293,12 @@ def test_partitioning(self):
             assert isinstance(partition, xugrid.UgridDataArray)
             assert partition.name == self.uda.name
 
+    def test_reindex_like(self):
+        back = self.uda.ugrid.reindex_like(self.uda)
+        assert isinstance(back, xugrid.UgridDataArray)
+        back = self.uda.ugrid.reindex_like(self.uda.ugrid.grid)
+        assert isinstance(back, xugrid.UgridDataArray)
+
     def test_crs(self):
         uda = self.uda
         crs = uda.ugrid.crs
@@ -637,6 +643,12 @@ def test_partitioning(self):
             assert "a" in partition
             assert "b" in partition
 
+    def test_reindex_like(self):
+        back = self.uds.ugrid.reindex_like(self.uds)
+        assert isinstance(back, xugrid.UgridDataset)
+        back = self.uds.ugrid.reindex_like(self.uds.ugrid.grid)
+        assert isinstance(back, xugrid.UgridDataset)
+
     def test_crs(self):
         uds = self.uds
         crs = uds.ugrid.crs
@@ -692,7 +704,7 @@ def test_total_bounds(self):
         assert self.uds.ugrid.total_bounds == (0.0, 0.0, 2.0, 2.0)
 
 
-class TestMultiToplogyUgridDataset:
+class TestMultiTopologyUgridDataset:
     @pytest.fixture(autouse=True)
     def setup(self):
         self.uds = xugrid.UgridDataset(grids=GRID())
@@ -732,6 +744,10 @@ def test_multi_topology_sel(self):
         # Ensure both grids are still present
         assert len(result.ugrid.grids) == 2
 
+    def test_reindex_like(self):
+        back = self.uds.ugrid.reindex_like(self.uds)
+        assert isinstance(back, xugrid.UgridDataset)
+
 
 def test_multiple_coordinates():
     grid = GRID()

xugrid/core/dataarray_accessor.py

@@ -7,10 +7,12 @@
 # from xugrid.plot.pyvista import to_pyvista_grid
 from xugrid.core.accessorbase import AbstractUgridAccessor
 from xugrid.core.utils import UncachedAccessor
-from xugrid.core.wrap import UgridDataArray
+from xugrid.core.wrap import UgridDataArray, UgridDataset
 from xugrid.plot.plot import _PlotMethods
 from xugrid.ugrid import connectivity
 from xugrid.ugrid.interpolate import laplace_interpolate
+from xugrid.ugrid.ugrid1d import Ugrid1d
+from xugrid.ugrid.ugrid2d import Ugrid2d
 from xugrid.ugrid.ugridbase import UgridType
 
 
@@ -418,6 +420,41 @@ def to_geodataframe(
         geometry = self.grid.to_shapely(dim)
         return gpd.GeoDataFrame(df, geometry=geometry, crs=self.grid.crs)
 
+    def reindex_like(
+        self,
+        other: Union[UgridType, UgridDataArray, UgridDataset],
+        tolerance: float = 0.0,
+    ):
+        """
+        Conform this object to match the topology of another object. The
+        topologies must be exactly equivalent: only the order of the nodes,
+        edges, and faces may differ.
+
+        Dimension names must match for equivalent topologies.
+
+        Parameters
+        ----------
+        other: Ugrid1d, Ugrid2d, UgridDataArray, UgridDataset
+        obj: DataArray or Dataset
+        tolerance: float, default value 0.0.
+            Maximum distance between inexact coordinate matches.
+
+        Returns
+        -------
+        reindexed: UgridDataArray
+        """
+        if isinstance(other, (Ugrid1d, Ugrid2d)):
+            other_grid = other
+        elif isinstance(other, (UgridDataArray, UgridDataset)):
+            other_grid = other.ugrid.grid
+        else:
+            raise TypeError(
+                "Expected Ugrid1d, Ugrid2d, UgridDataArray, or UgridDataset,"
+                f"received instead: {type(other).__name__}"
+            )
+        new_obj = self.grid.reindex_like(other_grid, obj=self.obj, tolerance=tolerance)
+        return UgridDataArray(new_obj, other_grid)
+
     def _binary_iterate(self, iterations: int, mask, value, border_value):
         if border_value == value:
             exterior = self.grid.exterior_faces

xugrid/core/dataset_accessor.py

@@ -6,7 +6,9 @@
 
 # from xugrid.plot.pyvista import to_pyvista_grid
 from xugrid.core.accessorbase import AbstractUgridAccessor
-from xugrid.core.wrap import UgridDataset
+from xugrid.core.wrap import UgridDataArray, UgridDataset
+from xugrid.ugrid.ugrid1d import Ugrid1d
+from xugrid.ugrid.ugrid2d import Ugrid2d
 from xugrid.ugrid.ugridbase import UgridType
 
 
@@ -378,7 +380,7 @@ def intersect_linestring(self, linestring) -> xr.Dataset:
 
     def to_dataset(self, optional_attributes: bool = False):
         """
-        Converts this UgridDataArray or UgridDataset into a standard
+        Converts this UgridDataset into a standard
         xarray.Dataset.
 
         The UGRID topology information is added as standard data variables.
@@ -554,3 +556,50 @@ def to_geodataframe(
             )
 
         return pd.concat(gdfs)
+
+    def reindex_like(
+        self,
+        other: Union[UgridType, UgridDataArray, UgridDataset],
+        tolerance: float = 0.0,
+    ):
+        """
+        Conform this object to match the topology of another object. The
+        topologies must be exactly equivalent: only the order of the nodes,
+        edges, and faces may differ.
+
+        Topologies are matched by name, and dimension names must match for
+        equivalent topologies.
+
+        Parameters
+        ----------
+        other: Ugrid1d, Ugrid2d, UgridDataArray, UgridDataset
+        obj: DataArray or Dataset
+        tolerance: float, default value 0.0.
+            Maximum distance between inexact coordinate matches.
+
+        Returns
+        -------
+        reindexed: UgridDataset
+        """
+        if isinstance(other, (Ugrid1d, Ugrid2d)):
+            other_grids = [other]
+        elif isinstance(other, (UgridDataArray, UgridDataset)):
+            other_grids = other.ugrid.grids
+        else:
+            raise TypeError(
+                "Expected Ugrid1d, Ugrid2d, UgridDataArray, or UgridDataset,"
+                f"received instead: {type(other).__name__}"
+            )
+        # Convert to dict to match by name
+        other_grids = {grid.name: grid for grid in other_grids}
+
+        new_grids = []
+        result = self.obj
+        for grid in self.grids:
+            other = other_grids.get(grid.name)
+            if other:
+                result = grid.reindex_like(other, obj=result, tolerance=tolerance)
+                new_grids.append(other)
+            else:
+                new_grids.append(grid)
+        return UgridDataset(result, new_grids)

xugrid/ugrid/connectivity.py

@@ -8,6 +8,33 @@
 from xugrid.constants import BoolArray, FloatArray, IntArray, IntDType, SparseMatrix
 
 
+def argsort_rows(array: np.ndarray) -> IntArray:
+    if array.ndim != 2:
+        raise ValueError(f"Array is not 2D, but has shape: {array.shape}")
+    dtype = [("f{i}".format(i=i), array.dtype) for i in range(array.shape[1])]
+    arr1d = array.view(dtype).flatten()
+    return np.argsort(arr1d)
+
+
+def index_like(xy_a: FloatArray, xy_b: FloatArray, tolerance: float):
+    """
+    Return the index that would transform xy_a into xy_b.
+    """
+    if xy_a.shape != xy_b.shape:
+        raise ValueError("coordinates do not match in shape")
+
+    sorter_a = argsort_rows(xy_a)
+    sorter_b = argsort_rows(xy_b)
+    if not np.allclose(xy_a[sorter_a], xy_b[sorter_b], rtol=0.0, atol=tolerance):
+        raise ValueError("coordinates are not identical after sorting")
+    #
+    #   sorter_a         inverse_b
+    # a --------> sorted ---------> b
+    #
+    inverse_b = np.argsort(sorter_b)
+    return sorter_a[inverse_b]
+
+
 class AdjacencyMatrix(NamedTuple):
     indices: IntArray
     indptr: IntArray