Merge branch 'main' into dim-sizes

.gitignore

@@ -72,6 +72,7 @@ docs/api
 docs/examples
 docs/examples-dev
 docs/sample_data
+docs/sg_execution_times.rst
 
 # PyBuilder
 target/

docs/api.rst

@@ -127,6 +127,7 @@ UgridDataArray or UgridDataset.
     UgridDataArrayAccessor.binary_erosion
     UgridDataArrayAccessor.connected_components
     UgridDataArrayAccessor.reverse_cuthill_mckee
+    UgridDataArrayAccessor.interpolate_na
     UgridDataArrayAccessor.laplace_interpolate
     UgridDataArrayAccessor.to_dataset
     UgridDataArrayAccessor.to_netcdf
@@ -193,13 +194,15 @@ UGRID1D Topology
     Ugrid1d.dims
     Ugrid1d.sizes
     Ugrid1d.attrs
+    Ugrid1d.coords
 
     Ugrid1d.n_node
     Ugrid1d.node_dimension
     Ugrid1d.node_coordinates
     Ugrid1d.set_node_coords
     Ugrid1d.assign_node_coords
     Ugrid1d.assign_edge_coords
+    Ugrid1d.get_coordinates
 
     Ugrid1d.n_edge
     Ugrid1d.edge_dimension
@@ -212,6 +215,7 @@ UGRID1D Topology
 
     Ugrid1d.node_edge_connectivity
     Ugrid1d.node_node_connectivity
+    Ugrid1d.get_connectivity_matrix
 
     Ugrid1d.copy
     Ugrid1d.rename
@@ -253,6 +257,7 @@ UGRID2D Topology
     Ugrid2d.dims
     Ugrid2d.sizes
     Ugrid2d.attrs
+    Ugrid2d.coords
 
     Ugrid2d.n_node
     Ugrid2d.node_dimension
@@ -261,6 +266,7 @@ UGRID2D Topology
     Ugrid2d.assign_node_coords
     Ugrid2d.assign_edge_coords
     Ugrid2d.assign_face_coords
+    Ugrid2d.get_coordinates
 
     Ugrid2d.n_edge
     Ugrid2d.edge_dimension
@@ -289,6 +295,7 @@ UGRID2D Topology
     Ugrid2d.edge_node_connectivity
     Ugrid2d.face_edge_connectivity
     Ugrid2d.face_face_connectivity
+    Ugrid2d.get_connectivity_matrix
 
     Ugrid2d.validate_edge_node_connectivity
 

docs/changelog.rst

@@ -16,7 +16,13 @@ Fixed
   linear weights within the full bounds of the source grid, rather than only
   within the centroids of the source grid. Previously, it would give no results
   beyond the centroids for structured to structured regridding, and it would
-  give nearest results (equal to :class:`CentroidLocatorRegridder`) otherwise.
+  give nearest results (equal to :class:`xugrid.CentroidLocatorRegridder`) otherwise.
+
+Added
+~~~~~
+
+- :meth:`UgridDataArrayAccessor.interpolate_na` has been added to fill missing
+  data. Currently, the only supported method is ``"nearest"``.
 
 Changed
 ~~~~~~~

tests/test_interpolate.py

@@ -50,3 +50,20 @@ def test_laplace_interpolate():
         con, data, use_weights=False, direct_solve=False
     )
     assert np.allclose(actual, expected)
+
+
+def test_nearest_interpolate():
+    x = np.array([0.0, 1.0, 2.0, 3.0, 4.0])
+    y = np.zeros_like(x)
+    coordinates = np.column_stack((x, y))
+    data = np.array([0.0, np.nan, np.nan, np.nan, 4.0])
+    actual = interpolate.nearest_interpolate(coordinates, data, np.inf)
+    assert np.allclose(actual, np.array([0.0, 0.0, 0.0, 4.0, 4.0]))
+
+    actual = interpolate.nearest_interpolate(coordinates, data, 1.1)
+    assert np.allclose(actual, np.array([0.0, 0.0, np.nan, 4.0, 4.0]), equal_nan=True)
+
+    with pytest.raises(ValueError, match="All values are NA."):
+        interpolate.nearest_interpolate(
+            coordinates, data=np.full_like(data, np.nan), max_distance=np.inf
+        )

tests/test_ugrid1d.py

@@ -105,6 +105,11 @@ def test_ugrid1d_properties():
     assert np.allclose(actual_coords, expected_coords)
     assert isinstance(grid.attrs, dict)
 
+    coords = grid.coords
+    assert isinstance(coords, dict)
+    assert np.array_equal(coords[grid.node_dimension], grid.node_coordinates)
+    assert np.array_equal(coords[grid.edge_dimension], grid.edge_coordinates)
+
 
 def test_ugrid1d_egde_bounds():
     grid = grid1d()
@@ -446,7 +451,7 @@ def test_contract_vertices():
     assert np.array_equal(new.edge_node_connectivity, [[0, 1]])
 
 
-def test_connectivity_matrix():
+def test_get_connectivity_matrix():
     xy = np.array(
         [
             [0.0, 0.0],
@@ -461,14 +466,37 @@ def test_connectivity_matrix():
         edge_node_connectivity=np.array([[0, 1], [1, 2]]),
     )
     with pytest.raises(ValueError, match="Expected network1d_nNodes; got: abc"):
-        grid.connectivity_matrix(dim="abc", xy_weights=True)
+        grid.get_connectivity_matrix(dim="abc", xy_weights=True)
 
-    connectivity = grid.connectivity_matrix(grid.node_dimension, True)
+    connectivity = grid.get_connectivity_matrix(grid.node_dimension, True)
     assert isinstance(connectivity, sparse.csr_matrix)
     assert np.allclose(connectivity.data, [1.5, 1.5, 0.75, 0.75])
     assert np.array_equal(connectivity.indices, [1, 0, 2, 1])
 
 
+def test_get_coordinates():
+    xy = np.array(
+        [
+            [0.0, 0.0],
+            [1.0, 0.0],
+            [3.0, 0.0],
+        ]
+    )
+    grid = xugrid.Ugrid1d(
+        node_x=xy[:, 0],
+        node_y=xy[:, 1],
+        fill_value=-1,
+        edge_node_connectivity=np.array([[0, 1], [1, 2]]),
+    )
+    with pytest.raises(
+        ValueError, match="Expected network1d_nNodes or network1d_nEdges; got: abc"
+    ):
+        grid.get_coordinates(dim="abc")
+
+    assert isinstance(grid.get_coordinates(grid.node_dimension), np.ndarray)
+    assert isinstance(grid.get_coordinates(grid.edge_dimension), np.ndarray)
+
+
 def test_equals():
     grid = grid1d()
     grid_copy = grid1d()

tests/test_ugrid2d.py

@@ -161,6 +161,11 @@ def test_ugrid2d_properties():
     assert are_nan[2:, -1:, :].all()
     assert not are_nan[:, :-1, :].any()
     assert isinstance(grid.attrs, dict)
+    coords = grid.coords
+    assert isinstance(coords, dict)
+    assert np.array_equal(coords[grid.node_dimension], grid.node_coordinates)
+    assert np.array_equal(coords[grid.edge_dimension], grid.edge_coordinates)
+    assert np.array_equal(coords[grid.face_dimension], grid.face_coordinates)
 
 
 def test_validate_edge_node_connectivity():
@@ -503,20 +508,33 @@ def test_connectivity_matrix():
     with pytest.raises(
         ValueError, match="Expected mesh2d_nNodes or mesh2d_nFaces; got: mesh2d_nEdges"
     ):
-        grid.connectivity_matrix(dim=grid.edge_dimension, xy_weights=False)
+        grid.get_connectivity_matrix(dim=grid.edge_dimension, xy_weights=False)
 
-    connectivity = grid.connectivity_matrix(grid.face_dimension, xy_weights=True)
+    connectivity = grid.get_connectivity_matrix(grid.face_dimension, xy_weights=True)
     assert isinstance(connectivity, sparse.csr_matrix)
     assert np.array_equal(connectivity.indices, [1, 2, 0, 3, 0, 3, 1, 2])
 
-    connectivity = grid.connectivity_matrix(grid.node_dimension, xy_weights=True)
+    connectivity = grid.get_connectivity_matrix(grid.node_dimension, xy_weights=True)
     assert isinstance(connectivity, sparse.csr_matrix)
     assert np.array_equal(
         connectivity.indices,
         [1, 3, 0, 2, 4, 1, 5, 0, 4, 6, 1, 3, 5, 6, 2, 4, 6, 3, 4, 5],
     )
 
 
+def test_get_coordinates():
+    grid = grid2d()
+    with pytest.raises(
+        ValueError,
+        match="Expected mesh2d_nNodes, mesh2d_nEdges, or mesh2d_nFaces; got: abc",
+    ):
+        grid.get_coordinates(dim="abc")
+
+    assert isinstance(grid.get_coordinates(grid.node_dimension), np.ndarray)
+    assert isinstance(grid.get_coordinates(grid.edge_dimension), np.ndarray)
+    assert isinstance(grid.get_coordinates(grid.face_dimension), np.ndarray)
+
+
 def test_voronoi_topology():
     grid = grid2d()
     vertices, faces, face_index = grid.voronoi_topology

tests/test_ugrid_dataset.py

@@ -69,6 +69,7 @@ def ugrid1d_ds():
     )
     ds = grid.to_dataset()
     ds["a1d"] = xr.DataArray([1.0, 2.0, 3.0], dims=[grid.node_dimension])
+    ds["b1d"] = xr.DataArray([1.0, 2.0], dims=[grid.edge_dimension])
     return xugrid.UgridDataset(ds)
 
 
@@ -1262,6 +1263,11 @@ def test_laplace_interpolate_facets():
     with pytest.raises(ValueError, match=msg):
         edge_uda.ugrid.laplace_interpolate(direct_solve=True)
 
+    for uda in (node_uda, edge_uda, face_uda):
+        actual = uda.ugrid.interpolate_na()
+        assert isinstance(actual, xugrid.UgridDataArray)
+        assert np.allclose(actual, 1.0)
+
 
 def test_laplace_interpolate_1d():
     uda = ugrid1d_ds()["a1d"]
@@ -1272,6 +1278,32 @@ def test_laplace_interpolate_1d():
     assert np.allclose(actual, 1.0)
 
 
+def test_interpolate_na_1d():
+    uda = ugrid1d_ds()["a1d"]
+    with pytest.raises(ValueError, match='"abc" is not a valid interpolator.'):
+        uda.ugrid.interpolate_na(method="abc")
+
+    # Node data
+    uda = ugrid1d_ds()["a1d"]
+    uda[:] = 1.0
+    uda[1] = np.nan
+    actual = uda.ugrid.interpolate_na()
+    assert isinstance(actual, xugrid.UgridDataArray)
+    assert np.allclose(actual, 1.0)
+
+    # Edge data
+    uda = ugrid1d_ds()["b1d"]
+    uda[:] = 1.0
+    uda[1] = np.nan
+    actual = uda.ugrid.interpolate_na()
+    assert isinstance(actual, xugrid.UgridDataArray)
+    assert np.allclose(actual, 1.0)
+
+    # Check max_distance
+    actual = uda.ugrid.interpolate_na(max_distance=0.5)
+    assert np.isnan(actual[1])
+
+
 def test_ugriddataset_wrap_twice(tmp_path):
     """
     in issue https://github.com/Deltares/xugrid/issues/208 wrapping a ds

xugrid/core/dataarray_accessor.py

@@ -1,4 +1,4 @@
-from typing import Dict, List, Sequence, Tuple, Union
+from typing import Dict, List, Optional, Sequence, Tuple, Union
 
 import numpy as np
 import scipy.sparse
@@ -10,7 +10,7 @@
 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.interpolate import laplace_interpolate, nearest_interpolate
 from xugrid.ugrid.ugrid1d import Ugrid1d
 from xugrid.ugrid.ugrid2d import Ugrid2d
 from xugrid.ugrid.ugridbase import UgridType
@@ -579,6 +579,43 @@ def reverse_cuthill_mckee(self):
             reordered_grid,
         )
 
+    def interpolate_na(
+        self,
+        method: str = "nearest",
+        max_distance: Optional[float] = None,
+    ):
+        """
+        Fill in NaNs by interpolating.
+
+        Parameters
+        ----------
+        method: str, default is "nearest"
+            Currently the only supported method.
+        max_distance: nonnegative float, optional.
+            Use ``None`` for no maximum distance.
+
+        Returns
+        -------
+        filled: UgridDataArray of floats
+        """
+
+        if method != "nearest":
+            raise ValueError(f'"{method}" is not a valid interpolator.')
+
+        if max_distance is None:
+            max_distance = np.inf
+
+        grid = self.grid
+        da = self.obj
+
+        filled = nearest_interpolate(
+            coordinates=grid.get_coordinates(dim=da.dims[0]),
+            data=da.to_numpy(),
+            max_distance=max_distance,
+        )
+        da_filled = da.copy(data=filled)
+        return UgridDataArray(da_filled, grid)
+
     def laplace_interpolate(
         self,
         xy_weights: bool = True,
@@ -590,7 +627,7 @@ def laplace_interpolate(
         maxiter: int = 500,
     ):
         """
-        Fill gaps in ``data`` (``np.nan`` values) using Laplace interpolation.
+        Fill in NaNs by using Laplace interpolation.
 
         This solves Laplace's equation where where there is no data, with data
         values functioning as fixed potential boundary conditions.
@@ -638,7 +675,7 @@ def laplace_interpolate(
         if da.dims[0] == grid.edge_dimension:
             raise ValueError("Laplace interpolation along edges is not allowed.")
 
-        connectivity = grid.connectivity_matrix(da.dims[0], xy_weights=xy_weights)
+        connectivity = grid.get_connectivity_matrix(da.dims[0], xy_weights=xy_weights)
         filled = laplace_interpolate(
             connectivity=connectivity,
             data=da.to_numpy(),

xugrid/ugrid/interpolate.py

@@ -6,6 +6,7 @@
 import numba as nb
 import numpy as np
 from scipy import sparse
+from scipy.spatial import KDTree
 
 from xugrid.constants import FloatArray, IntArray
 from xugrid.core.sparse import MatrixCSR, columns_and_values, nzrange, row_slice
@@ -306,3 +307,33 @@ def laplace_interpolate(
             warnings.warn(f"Failed to converge after {maxiter} iterations")
 
     return x
+
+
+def nearest_interpolate(
+    coordinates: FloatArray,
+    data: FloatArray,
+    max_distance: float,
+) -> FloatArray:
+    isnull = np.isnan(data)
+    if isnull.all():
+        raise ValueError("All values are NA.")
+
+    i_source = np.flatnonzero(~isnull)
+    i_target = np.flatnonzero(isnull)
+    source_coordinates = coordinates[i_source]
+    target_coordinates = coordinates[i_target]
+    # Locate the nearest notnull for each null value.
+    tree = KDTree(source_coordinates)
+    _, index = tree.query(
+        target_coordinates, distance_upper_bound=max_distance, workers=-1
+    )
+    # Remove entries beyond max distance, returned by .query as self.n.
+    keep = index < len(source_coordinates)
+    index = index[keep]
+    i_target = i_target[keep]
+    # index contains an index of the target coordinates to the source
+    # coordinates, not the direct index into the data, so we need an additional
+    # indexing step.
+    out = data.copy()
+    out[i_target] = data[i_source[index]]
+    return out