Fix stale flow-direction caches (#128)

* fix stream_order cache

* other cache fixes

Author: DirkEilander

pyflwdir/flwdir.py

@@ -291,7 +291,7 @@ def main_upstream(self, uparea: np.ndarray | None = None) -> np.ndarray:
         idxs_us_main = core.main_upstream(
             idxs_ds=self.idxs_ds, uparea=self._check_data(uparea, "uparea"), mv=self._mv
         )
-        if self.cache:
+        if self.cache and uparea is None:
             self._cached.update(idxs_us_main=idxs_us_main)
         return idxs_us_main
 
@@ -313,10 +313,11 @@ def add_pits(
         # add pits
         self.idxs_ds[idxs1] = idxs1
         self._pit = np.unique(np.concatenate([self.idxs_pit, idxs1]))
-        # reset order, nnodes and upstream cell indices
+        # Reset traversal state and all values derived from the flow topology.
         self._seq = None
         self._nnodes = None
-        self._idxs_us_main = None
+        for key in ("rank", "strord", "idxs_us_main", "distnc"):
+            self._cached.pop(key, None)
 
     def repair_loops(self) -> None:
         """Repair loops by setting a pit at every cell which does not drain to a pit."""
@@ -593,11 +594,11 @@ def stream_order(
         """
         mask = self._check_data(mask, "mask", optional=True)
         if type.lower() == "strahler":
-            if "strord" in self._cached:
+            if mask is None and "strord" in self._cached:
                 strord = self._cached["strord"]
             else:
                 strord = streams.strahler_order(self.idxs_ds, self.idxs_seq, mask=mask)
-                if self.cache:
+                if self.cache and mask is None:
                     self._cached.update(strord=strord)
         elif type.lower() == "classic":
             strord = streams.stream_order(

pyflwdir/pyflwdir.py

@@ -343,6 +343,8 @@ def set_transform(self, transform: Affine, latlon: bool = False) -> None:
                 raise ValueError("Invalid transform.")
         self.transform = transform
         self.latlon = latlon
+        for key in ("area", "distnc", "idxs_us_main"):
+            self._cached.pop(key, None)
 
     ### WRITE / EXPORT ###
 

tests/test_pyflwdir.py

@@ -9,7 +9,7 @@
 from affine import Affine
 
 import pyflwdir
-from pyflwdir import core
+from pyflwdir import core, streams
 from pyflwdir.pyflwdir import FlwdirRaster, _get_idxs_dtype
 
 pyflwdir_module = importlib.import_module("pyflwdir.pyflwdir")
@@ -381,6 +381,129 @@ def test_streams(flw_real, flwdir_real_rank):
     assert np.all(data_smooth1 == data)
 
 
+@pytest.mark.integration
+@pytest.mark.parametrize("raster", [False, True])
+def test_stream_order_mask_cache(flw_real, raster):
+    if raster:
+        flw = FlwdirRaster(
+            flw_real.idxs_ds.copy(),
+            flw_real.shape,
+            "d8",
+            idxs_pit=flw_real.idxs_pit.copy(),
+            cache=True,
+        )
+    else:
+        flw = pyflwdir.Flwdir(
+            flw_real.idxs_ds.copy(),
+            idxs_pit=flw_real.idxs_pit.copy(),
+            cache=True,
+        )
+
+    all_streams = flw.mask.reshape(flw.shape)
+    pit_streams = np.zeros(flw.shape, dtype=bool)
+    pit_streams.flat[flw.idxs_pit] = True
+
+    all_order = flw.stream_order(mask=all_streams)
+    pit_order = flw.stream_order(mask=pit_streams)
+    expected = streams.strahler_order(
+        flw.idxs_ds, flw.idxs_seq, mask=pit_streams.ravel()
+    ).reshape(flw.shape)
+
+    assert not np.array_equal(all_order, pit_order)
+    assert np.array_equal(pit_order, expected)
+    assert np.all(pit_order[~pit_streams] == 0)
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize("raster", [False, True])
+def test_add_pits_invalidates_cache(raster):
+    idxs_ds = np.array([0, 0, 1, 1, core._mv], dtype=np.int32)
+    idxs_pit = np.array([0], dtype=np.int32)
+    if raster:
+        flw = FlwdirRaster(
+            idxs_ds.copy(), (1, 5), "d8", idxs_pit=idxs_pit.copy(), cache=True
+        )
+    else:
+        flw = pyflwdir.Flwdir(idxs_ds.copy(), idxs_pit=idxs_pit.copy(), cache=True)
+
+    old_rank = flw.rank.copy()
+    old_strord = flw.stream_order().copy()
+    old_idxs_us_main = flw.idxs_us_main.copy()
+    if raster:
+        old_distnc = flw.distnc.copy()
+
+    flw.add_pits(idxs=np.array([2]))
+
+    assert "rank" not in flw._cached
+    assert "strord" not in flw._cached
+    assert "idxs_us_main" not in flw._cached
+    assert np.array_equal(
+        flw.rank, core.rank(flw.idxs_ds, mv=flw._mv)[0].reshape(flw.shape)
+    )
+    assert not np.array_equal(flw.rank, old_rank)
+    assert not np.array_equal(flw.stream_order(), old_strord)
+    assert flw.idxs_us_main[1] == 3
+    assert not np.array_equal(flw.idxs_us_main, old_idxs_us_main)
+    if raster:
+        assert "distnc" not in flw._cached
+        assert not np.array_equal(flw.distnc, old_distnc)
+
+
+@pytest.mark.unit
+def test_repair_loops_invalidates_topology_cache():
+    idxs_ds = np.array([0, 2, 1], dtype=np.int32)
+    flw = pyflwdir.Flwdir(idxs_ds, idxs_pit=np.array([0], dtype=np.int32), cache=True)
+    old_rank = flw.rank.copy()
+    old_strord = flw.stream_order().copy()
+    old_idxs_us_main = flw.idxs_us_main.copy()
+
+    flw.repair_loops()
+
+    assert flw.isvalid
+    assert "strord" not in flw._cached
+    assert "idxs_us_main" not in flw._cached
+    assert np.array_equal(flw.rank, core.rank(flw.idxs_ds, mv=flw._mv)[0])
+    assert not np.array_equal(flw.rank, old_rank)
+    assert not np.array_equal(flw.stream_order(), old_strord)
+    assert not np.array_equal(flw.idxs_us_main, old_idxs_us_main)
+
+
+@pytest.mark.unit
+def test_set_transform_invalidates_geometry_cache():
+    idxs_ds = np.array([0, 0, 1, 1, core._mv], dtype=np.int32)
+    flw = FlwdirRaster(
+        idxs_ds, (1, 5), "d8", idxs_pit=np.array([0], dtype=np.int32), cache=True
+    )
+    old_area = flw.area.copy()
+    old_distnc = flw.distnc.copy()
+    old_idxs_us_main = flw.idxs_us_main.copy()
+
+    flw.set_transform(Affine.scale(2), latlon=False)
+
+    assert "area" not in flw._cached
+    assert "distnc" not in flw._cached
+    assert "idxs_us_main" not in flw._cached
+    mask = flw.mask.reshape(flw.shape)
+    assert np.all(flw.area[mask] == 4)
+    assert not np.array_equal(flw.area, old_area)
+    assert np.all(flw.distnc[mask] == 2 * old_distnc[mask])
+    assert np.array_equal(flw.idxs_us_main, old_idxs_us_main)
+
+
+@pytest.mark.unit
+def test_main_upstream_custom_area_does_not_replace_default_cache():
+    idxs_ds = np.array([0, 0, 1, 1, core._mv], dtype=np.int32)
+    flw = pyflwdir.Flwdir(idxs_ds, idxs_pit=np.array([0], dtype=np.int32), cache=True)
+    default_idxs_us_main = flw.idxs_us_main.copy()
+    custom_uparea = np.array([1, 1, 1, 2, 0], dtype=np.float32)
+
+    custom_idxs_us_main = flw.main_upstream(uparea=custom_uparea)
+
+    assert custom_idxs_us_main[1] == 3
+    assert default_idxs_us_main[1] == 2
+    assert np.array_equal(flw.idxs_us_main, default_idxs_us_main)
+
+
 @pytest.mark.integration
 def test_upscale(flw_real, nextxy_real):
     flw1, idxs_out = flw_real.upscale(5, method="dmm")  # single method