Supply boundary conditions to fillsinks and handle NaNs properly (Top…

…oToolbox#70)

Resolves TopoToolbox#69.

CMakeLists.txt: Set GIT_TAG to main to access the new boundary
conditions API in libtopotoolbox. This should be changed to 2024-W43
before this is merged but after the next weekly release of libtopotoolbox

src/lib/grid.cpp: `wrap_fillsinks` adds a `bc` argument, which it
supplies to libtopotoolbox's `fillsinks` directly.

src/topotoolbox/grid_object.py: The boundary conditions are supplied
to `GridObject.fillsinks` as an optional argument. If `bc` is not
supplied, then the default is to fix the DEM to its value on the
exterior and fill all interior pixels that are not NaNs. NaNs are
treated as if they were known sinks by temporarily setting the DEM to
-infinity and fixing that pixel as if it were a boundary pixel.

If the user passes in their own boundary condition array, we assume
that they have already handled NaNs according to their own
preferences, and we do not treat them in any particular way.

src/topotoolbox/flow_object.py adds the same boundary condition logic
as in `GridObjects.fillsinks`. The user can optionally pass a `bc`
parameter to the `FlowObject` constructor. We could use this in the
future to implement more complex controls on flow routing, so I am not
too worried about adding an extra optional argument here.

This has been informally tested against the TopoToolbox snapshot
tests, but those have not yet been implemented in pytopotoolbox.

CMakeLists.txt

@@ -9,7 +9,7 @@ include(FetchContent)
 FetchContent_Declare(
   topotoolbox
   GIT_REPOSITORY https://github.com/TopoToolbox/libtopotoolbox.git
-  GIT_TAG 2024-W40
+  GIT_TAG main
 )
 FetchContent_MakeAvailable(topotoolbox)
 

src/lib/grid.cpp

@@ -19,15 +19,17 @@ namespace py = pybind11;
 //   dem: A NumPy array representing the digital elevation model.
 //   dims: A tuple containing the number of rows and columns.
 
-void wrap_fillsinks(py::array_t<float> output, py::array_t<float> dem, 
+void wrap_fillsinks(py::array_t<float> output, py::array_t<float> dem,
+                    py::array_t<uint8_t> bc,
                     std::tuple<ptrdiff_t, ptrdiff_t> dims){
 
     float *output_ptr = output.mutable_data();
     float *dem_ptr = dem.mutable_data();
+    uint8_t *bc_ptr = bc.mutable_data();
 
     std::array<ptrdiff_t, 2> dims_array = {std::get<0>(dims), std::get<1>(dims)};
     ptrdiff_t *dims_ptr = dims_array.data();
-    fillsinks(output_ptr, dem_ptr, dims_ptr);
+    fillsinks(output_ptr, dem_ptr, bc_ptr, dims_ptr);
 }
 
 // wrap_identifyflats: 

src/topotoolbox/flow_object.py

@@ -15,14 +15,21 @@ class FlowObject():
     digital elevation model (DEM).
     """
 
-    def __init__(self, grid: GridObject):
+    def __init__(self, grid: GridObject,
+                 bc: np.ndarray | GridObject | None = None):
         """The constructor for the FlowObject. Takes a GridObject as input,
         computes flow direction information and saves them as an FlowObject.
 
         Parameters
         ----------
         grid : GridObject
             The GridObject that will be the basis of the computation.
+        bc : ndarray or GridObject, optional
+            Boundary conditions for sink filling. `bc` should be an array
+            of np.uint8 that matches the shape of the DEM. Values of 1
+            indicate pixels that should be fixed to their values in the
+            original DEM and values of 0 indicate pixels that should be
+            filled.
 
         Notes
         -----
@@ -33,7 +40,28 @@ def __init__(self, grid: GridObject):
         dem = grid.z
 
         filled_dem = np.zeros_like(dem, dtype=np.float32, order='F')
-        _grid.fillsinks(filled_dem, dem, dims)
+        restore_nans = False
+        if bc is None:
+            bc = np.ones_like(dem, dtype=np.uint8)
+            bc[1:-1, 1:-1] = 0  # Set interior pixels to 0
+
+            nans = np.isnan(dem)
+            dem[nans] = -np.inf
+            bc[nans] = 1
+            restore_nans = True
+
+        if bc.shape != dims:
+            err = ("The shape of the provided boundary conditions does not "
+                   f"match the shape of the DEM. {dims}")
+            raise ValueError(err)from None
+
+        if isinstance(bc, GridObject):
+            bc = bc.z
+
+        _grid.fillsinks(filled_dem, dem, bc, dims)
+
+        if restore_nans:
+            filled_dem[nans] = np.nan
 
         flats = np.zeros_like(dem, dtype=np.int32, order='F')
         _grid.identifyflats(flats, filled_dem, dims)

src/topotoolbox/grid_object.py

@@ -38,19 +38,52 @@ def __init__(self) -> None:
         self.transform = None
         self.crs = None
 
-    def fillsinks(self) -> 'GridObject':
+    def fillsinks(self,
+                  bc: 'np.ndarray | GridObject | None' = None) -> 'GridObject':
         """Fill sinks in the digital elevation model (DEM).
 
+        Parameters
+        ----------
+        bc : ndarray or GridObject, optional
+            Boundary conditions for sink filling. `bc` should be an array
+            of np.uint8 that matches the shape of the DEM. Values of 1
+            indicate pixels that should be fixed to their values in the
+            original DEM and values of 0 indicate pixels that should be
+            filled.
+
         Returns
         -------
         GridObject
             The filled DEM.
+
         """
 
         dem = self.z.astype(np.float32, order='F')
         output = np.zeros_like(dem)
 
-        _grid.fillsinks(output, dem, self.shape)
+        restore_nans = False
+
+        if bc is None:
+            bc = np.ones_like(dem, dtype=np.uint8)
+            bc[1:-1, 1:-1] = 0  # Set interior pixels to 0
+
+            nans = np.isnan(dem)
+            dem[nans] = -np.inf
+            bc[nans] = 1  # Set NaNs to 1
+            restore_nans = True
+
+        if bc.shape != self.shape:
+            err = ("The shape of the provided boundary conditions does not "
+                   f"match the shape of the DEM. {self.shape}")
+            raise ValueError(err)from None
+
+        if isinstance(bc, GridObject):
+            bc = bc.z
+
+        _grid.fillsinks(output, dem, bc, self.shape)
+
+        if restore_nans:
+            output[nans] = np.nan
 
         result = copy.copy(self)
         result.z = output