further distance tests for 2D and ND

Author: Oisin-M

src/earthkit/hydro/distance.py

@@ -77,11 +77,13 @@ def _upstream_distance_ND(river_network, field, grouping, mv):
     downstream_check_shorter[0] = river_network.downstream_nodes[
         downstream_check_shorter[0]
     ]
-    field[replace_upstream_indices] = np.max(field[downstream_replace_upstream]) + 1
-    np.minimum.at(
-        field, replace_upstream_indices, field[downstream_replace_upstream] + 1
-    )
-    np.minimum.at(field, check_shorter_indices, field[downstream_check_shorter] + 1)
+    if replace_upstream_indices.size != 0:
+        field[replace_upstream_indices] = np.max(field[downstream_replace_upstream]) + 1
+        np.minimum.at(
+            field, replace_upstream_indices, field[downstream_replace_upstream] + 1
+        )
+    if check_shorter_indices.size != 0:
+        np.minimum.at(field, check_shorter_indices, field[downstream_check_shorter] + 1)
 
 
 def _downstream_distance_2D(river_network, field, grouping, mv):
@@ -125,8 +127,12 @@ def _downstream_distance_ND(river_network, field, grouping, mv):
     downstream_check_shorter[0] = river_network.downstream_nodes[
         downstream_check_shorter[0]
     ]
-    field[downstream_replace_downstream] = np.max(field[replace_downstream_indices]) + 1
-    np.minimum.at(
-        field, downstream_replace_downstream, field[replace_downstream_indices] + 1
-    )
-    np.minimum.at(field, downstream_check_shorter, field[check_shorter_indices] + 1)
+    if downstream_replace_downstream.size != 0:
+        field[downstream_replace_downstream] = (
+            np.max(field[replace_downstream_indices]) + 1
+        )
+        np.minimum.at(
+            field, downstream_replace_downstream, field[replace_downstream_indices] + 1
+        )
+    if downstream_check_shorter.size != 0:
+        np.minimum.at(field, downstream_check_shorter, field[check_shorter_indices] + 1)

tests/test_river_network.py

@@ -333,3 +333,43 @@ def test_distance(reader, map_name, distance_query, distance):
     print(dist)
     print(distance)
     np.testing.assert_array_equal(dist, distance)
+
+
+@parametrize(
+    "reader,map_name,distance_query,distance",
+    [
+        ("d8_ldd", d8_ldd_1, distance_query_field_1, distance_1),
+        ("cama_downxy", cama_downxy_1, distance_query_field_1, distance_1),
+        ("cama_nextxy", cama_nextxy_1, distance_query_field_1, distance_1),
+    ],
+)
+def test_distance_2d(reader, map_name, distance_query, distance):
+    network = read_network(reader, map_name)
+    field = np.zeros(network.mask.shape, dtype="int") - 1
+    field[network.mask] = distance_query
+    network_dist = ekh.compute_distance(network, field)
+    print(distance)
+    print(network_dist)
+    np.testing.assert_array_equal(network_dist[network.mask], distance)
+    np.testing.assert_array_equal(network_dist[~network.mask], -1)
+
+
+@parametrize(
+    "reader,map_name,distance_query,distance",
+    [
+        ("d8_ldd", d8_ldd_1, distance_query_field_1, distance_1),
+        ("cama_downxy", cama_downxy_1, distance_query_field_1, distance_1),
+        ("cama_nextxy", cama_nextxy_1, distance_query_field_1, distance_1),
+    ],
+)
+def test_distance_Nd(reader, map_name, distance_query, distance):
+    network = read_network(reader, map_name)
+    field = np.zeros(network.mask.shape, dtype="int") - 1
+    field[network.mask] = distance_query
+    field = np.stack([field, field], axis=0)
+    network_dist = ekh.compute_distance(network, field)
+    distance = np.stack([distance, distance], axis=0)
+    print(distance)
+    print(network_dist)
+    np.testing.assert_array_equal(network_dist[..., network.mask], distance)
+    np.testing.assert_array_equal(network_dist[..., ~network.mask], -1)