Repository: Deltares/pyflwdir

Issue:
No implementation of function Function(<built-in function heappush>) found for signature
I'm trying to get started with pyflwdir -- it looks very powerful. But the code in the quickstart at https://deltares.github.io/pyflwdir/latest/quickstart.html fails when executed in a newly created Python 3.11.6 virtualenv (on Linux, Debian sid, if that matters) after pip installing pyflwdir and rasterio:
 (and sorry for the formatting -- I'm using "code <>" but github's markdown seems to be treating python interpreter output like an email message instead of code):

`
>>> import rasterio
>>> with rasterio.open("n35_w107_1arc_v3.tif", "r") as src:
...     elevtn = src.read(1)
...     nodata = src.nodata
...     transform = src.transform
...     crs = src.crs
... 
>>> import pyflwdir
>>> flw = pyflwdir.from_dem(
...     data=elevtn,
...     nodata=src.nodata,
...     transform=transform,
...     latlon=crs.is_geographic,
... )
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/akkana/Data/demnames/dem-full-env/lib/python3.11/site-packages/pyflwdir/pyflwdir.py", line 97, in from_dem
    d8 = dem.fill_depressions(
         ^^^^^^^^^^^^^^^^^^^^^
  File "/home/akkana/Data/demnames/dem-full-env/lib/python3.11/site-packages/numba/core/dispatcher.py", line 468, in _compile_for_args
    error_rewrite(e, 'typing')
  File "/home/akkana/Data/demnames/dem-full-env/lib/python3.11/site-packages/numba/core/dispatcher.py", line 409, in error_rewrite
    raise e.with_traceback(None)
numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<built-in function heappush>) found for signature:
 
 >>> heappush(list(Tuple(int16, uint8, uint32, uint32))<iv=None>, Tuple(int64, uint8, uint32, uint32))
 
There are 2 candidate implementations:
  - Of which 2 did not match due to:
  Overload in function 'heappush': File: numba/cpython/heapq.py: Line 150.
    With argument(s): '(list(Tuple(int16, uint8, uint32, uint32))<iv=None>, Tuple(int64, uint8, uint32, uint32))':
   Rejected as the implementation raised a specific error:
     TypingError: heap type must be the same as item type
  raised from /home/akkana/Data/demnames/dem-full-env/lib/python3.11/site-packages/numba/cpython/heapq.py:119

During: resolving callee type: Function(<built-in function heappush>)
During: typing of call at /home/akkana/Data/demnames/dem-full-env/lib/python3.11/site-packages/pyflwdir/dem.py (130)


File "dem-full-env/lib/python3.11/site-packages/pyflwdir/dem.py", line 130:
def fill_depressions(
    <source elided>
            if ~queued[r, c]:  # add to queue
                heapq.heappush(q, (z1, np.uint8(0), np.uint32(r), np.uint32(c)))
                ^
`

This error was mentioned in Issue #10 and there was a request to try a more recent numba, but that was in 2021 and I see this with the numba currently in pip.

The elevation data I'm using is SRTM downloaded four years ago -- maybe there's a newer format that pyflwdir is expecting?

Discussion:

Patch:
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 99d12cb..92533a9 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -2,6 +2,11 @@
 Change Log
 ###########
 
+in development
+**************
+
+* Fixed support for interger type DEMs in `from_dem` and `dem.fill_depressions` methods
+
 0.5.8  (06-Oct-2023)
 ********************
 
diff --git a/pyflwdir/dem.py b/pyflwdir/dem.py
index dd67b6b..67c3f00 100644
--- a/pyflwdir/dem.py
+++ b/pyflwdir/dem.py
@@ -70,7 +70,7 @@ def fill_depressions(
     done = np.isnan(elevtn) if np.isnan(nodata) else elevtn == nodata
     d8 = np.where(done, np.uint8(247), np.uint8(0))
     if connectivity not in [4, 8]:
-        raise ValueError(f'"connectivity" should either be 4 or 8')
+        raise ValueError('"connectivity" should either be 4 or 8')
     # pfff.. numba does not allow creation of numpy bool arrays using normal methods
     struct = np.array([bool(1) for s in range(9)]).reshape((3, 3))
     if connectivity == 4:
@@ -91,10 +91,15 @@ def fill_depressions(
 
     # queue contains (elevation, boundary, row, col)
     # boundary is included to favor non-boundary cells over boundary cells with same elevation
-    q = [(elevtn[0, 0], np.uint8(1), np.uint32(0), np.uint32(0)) for _ in range(0)]
+    q = [
+        (np.float32(elevtn[0, 0]), np.uint8(1), np.uint32(0), np.uint32(0))
+        for _ in range(0)
+    ]
     heapq.heapify(q)
     for r, c in zip(*np.where(queued)):
-        heapq.heappush(q, (elevtn[r, c], np.uint8(1), np.uint32(r), np.uint32(c)))
+        heapq.heappush(
+            q, (np.float32(elevtn[r, c]), np.uint8(1), np.uint32(r), np.uint32(c))
+        )
     # restrict queue to global edge mimimum (single outlet)
     if outlets == "min":
         q = [heapq.heappop(q)]
@@ -115,7 +120,9 @@ def fill_depressions(
             dz = z0 - z1  # local depression if dz > 0
             if max_depth >= 0:  # if positive max_depth: don't fill when dz > max_depth
                 if dz >= max_depth:
-                    heapq.heappush(q, (z1, np.uint8(0), np.uint32(r), np.uint32(c)))
+                    heapq.heappush(
+                        q, (np.float32(z1), np.uint8(0), np.uint32(r), np.uint32(c))
+                    )
                     queued[r, c] = True
                     for dr, dc in zip(drs, dcs):  # (re)visit neighbors
                         done[r + dr, c + dc] = False
@@ -127,7 +134,9 @@ def fill_depressions(
                 delv[r, c] = dz
                 z1 += dz
             if ~queued[r, c]:  # add to queue
-                heapq.heappush(q, (z1, np.uint8(0), np.uint32(r), np.uint32(c)))
+                heapq.heappush(
+                    q, (np.float32(z1), np.uint8(0), np.uint32(r), np.uint32(c))
+                )
                 queued[r, c] = True
             done[r, c] = True
             d8[r, c] = core_d8._us[dr + 1, dc + 1]


Test patch:
diff --git a/tests/test_dem.py b/tests/test_dem.py
index 1fc3fec..aa1b317 100644
--- a/tests/test_dem.py
+++ b/tests/test_dem.py
@@ -10,7 +10,8 @@ parsed, flwdir = test_data[0]
 idxs_ds, idxs_pit, seq, rank, mv = parsed
 
 
-def test_from_dem():
+@pytest.mark.parametrize("dtype", [np.float32, np.int32])
+def test_from_dem(dtype):
     # example from Wang & Lui (2015)
     a = np.array(
         [
@@ -20,7 +21,7 @@ def test_from_dem():
             [16, 17, 8, 16, 15, 7, 5],
             [19, 18, 19, 18, 17, 15, 14],
         ],
-        dtype=np.float32,
+        dtype=dtype,
     )
     # NOTE: compared to paper same a_filled, but difference
     # in flowdir because row instead of col first ..
@@ -45,6 +46,15 @@ def test_from_dem():
     a2[0, 5] = 12
     a_filled = dem.fill_depressions(a, outlets="min")[0]
     assert np.all(a2 == a_filled)
+    # test with nodata values
+    a2 = a.copy()
+    a2[3, 5:] = -9999
+    _d8 = dem.fill_depressions(a2)[1]
+    assert np.all(_d8[3, 5:] == 247)
+    assert _d8[2, 4] == 0
+    # test max depth
+    _a = dem.fill_depressions(a, max_depth=2)[0]
+    assert np.all(a == _a)
     # test with 4-connectivity
     a2 = np.array(
         [
@@ -59,14 +69,6 @@ def test_from_dem():
     a_filled, _d8 = dem.fill_depressions(a, connectivity=4)
     assert np.all(a_filled == a2)
     assert np.all(np.isin(np.unique(_d8), [0, 1, 4, 16, 64]))
-    # test with nodata values
-    a[3, 5:] = -9999
-    _d8 = dem.fill_depressions(a)[1]
-    assert np.all(_d8[3, 5:] == 247)
-    assert _d8[2, 4] == 0
-    # test max depth
-    _a = dem.fill_depressions(a, max_depth=2)[0]
-    assert np.all(a == _a)
 
 
 def test_dem_adjust():