Improve testing coverage for hydrological functions with multi-backend support

tests/_test_inputs/accumulation.py

@@ -259,6 +259,22 @@
     dtype=float,
 )
 
+upstream_metric_var_1c = np.array(
+    [0.0, 0.0, 0.0, 0.0, 0.0,
+     3.24, 3.61, 26.08222222, 7.29, 0.0,
+     2.16222222, 7.10888889, 18.17061728, 11.12, 0.0,
+     0.0, 21.772475, 3.14888889, 2.7225, 0.0],
+    dtype=float,
+)
+
+upstream_metric_std_1c = np.array(
+    [0.0, 0.0, 0.0, 0.0, 0.0,
+     1.8, 1.9, 5.1070757, 2.7, 0.0,
+     1.47044967, 2.66624997, 4.2627007, 3.3346664, 0.0,
+     0.0, 4.66609848, 1.77451089, 1.65, 0.0],
+    dtype=float,
+)
+
 # 1d: bool field input
 input_field_1d = np.array(
     [
@@ -481,6 +497,22 @@
     dtype=float,
 )
 
+upstream_metric_var_1e = np.array(
+    [0.0, 0.0, 0.0, 0.0, 0.0,
+     3.24, 3.61, np.nan, 7.29, 0.0,
+     2.16222222, 7.10888889, np.nan, 11.12, 0.0,
+     0.0, np.nan, np.nan, np.nan, 0.0],
+    dtype=float,
+)
+
+upstream_metric_std_1e = np.array(
+    [0.0, 0.0, 0.0, 0.0, 0.0,
+     1.8, 1.9, np.nan, 2.7, 0.0,
+     1.47044967, 2.66624997, np.nan, 3.3346664, 0.0,
+     0.0, np.nan, np.nan, np.nan, 0.0],
+    dtype=float,
+)
+
 # 1f: missing float field input with mv=0
 input_field_1f = np.nan_to_num(input_field_1e, nan=0)
 
@@ -655,3 +687,70 @@
 upstream_metric_sum_2g = np.array(
     [-1, 2, -1, 4, 5, 11, -1, -1, 9, 10, -1, -1, -1, -1, 15, 31], dtype=int
 )
+
+# DOWNSTREAM ACCUMULATION RESULTS
+
+# Downstream sum for network 1, field 1c
+downstream_metric_sum_1c = np.array(
+    [7.7, 20.9, 18.8, 10.0, 9.9, 6.2, 13.6, 13.2, 9.8,
+     9.1, 1.1, 2.5, 4.3, 7.6, 3.1, -11.0, -2.1, 3.1,
+     7.5, 8.6],
+    dtype=float,
+)
+
+# Downstream sum for network 1, field 1e (with NaN values)
+downstream_metric_sum_1e = np.array(
+    [7.7, 20.9, np.nan, np.nan, 9.9, 6.2, 13.6, np.nan, 9.8,
+     9.1, 1.1, 2.5, 4.3, 7.6, 3.1, -11.0, -2.1, 3.1,
+     np.nan, np.nan],
+    dtype=float,
+)
+
+# Downstream sum for network 2, field 2g (with missing value -1)
+downstream_metric_sum_2g = np.array(
+    [1.0, 53.0, 51.0, 4.0, 36.0, 31.0, 32.0, 52.0, 23.0, 24.0, 25.0, 44.0, 0.0,
+     14.0, 75.0, 60.0],
+    dtype=float,
+)
+
+# Downstream min for network 1, field 1c
+downstream_metric_min_1c = np.array(
+    [-2.1, -2.1, -2.1, -3.2, -2.1, -2.1, -2.1, -2.1, -2.1, -2.1, -2.1,
+     -2.1, -2.1, -2.1, -4.5, -8.9, -2.1, -2.1, -2.1, -2.1],
+    dtype=float,
+)
+
+# Downstream min for network 1, field 1e (with NaN values)
+downstream_metric_min_1e = np.array(
+    [-2.1, -2.1, np.nan, np.nan, -2.1, -2.1, -2.1, np.nan, -2.1, -2.1, -2.1,
+     -2.1, -2.1, -2.1, -4.5, -8.9, -2.1, -2.1, np.nan, np.nan],
+    dtype=float,
+)
+
+# Downstream max for network 1, field 1c
+downstream_metric_max_1c = np.array(
+    [5.1, 11.1, 8.9, 8.9, 6.4, 5.1, 11.1, 8.9, 6.4, 6.4, 3.2,
+     4.6, 6.4, 6.4, 6.4, -2.1, -2.1, 5.2, 5.2, 5.2],
+    dtype=float,
+)
+
+# Downstream max for network 1, field 1e (with NaN values)
+downstream_metric_max_1e = np.array(
+    [5.1, 11.1, np.nan, np.nan, 6.4, 5.1, 11.1, np.nan, 6.4, 6.4, 3.2,
+     4.6, 6.4, 6.4, 6.4, -2.1, -2.1, 5.2, np.nan, np.nan],
+    dtype=float,
+)
+
+# Downstream mean for network 1, field 1c
+downstream_metric_mean_1c = np.array(
+    [1.925, 5.225, 4.7, 2.5, 2.475, 2.06666667, 4.53333333, 4.4, 3.26666667, 2.275,
+     0.55, 1.25, 2.15, 2.53333333, 0.775, -5.5, -2.1, 1.55, 2.5, 2.15],
+    dtype=float,
+)
+
+# Downstream mean for network 1, field 1e (with NaN values)
+downstream_metric_mean_1e = np.array(
+    [1.925, 5.225, np.nan, np.nan, 2.475, 2.06666667, 4.53333333, np.nan, 3.26666667, 2.275,
+     0.55, 1.25, 2.15, 2.53333333, 0.775, -5.5, -2.1, 1.55, np.nan, np.nan],
+    dtype=float,
+)

tests/_test_inputs/catchment.py

@@ -48,3 +48,10 @@
 
 
 catchment_2 = np.array([4, 2, 2, np.nan, 2, 2, 2, 2, 2, 2, 2, 2, 4, 2, 2, 2]) - 1
+
+# Catchment aggregation results for catchment_query_field_1 with input_field_1c
+# Locations: [8, 12, 13, 11, 10]
+catchment_sum_1c = np.array([5.6, 23.6, 0.3, 23.0, 9.8])
+catchment_mean_1c = np.array([2.8, 2.62222222, 0.1, 7.66666667, 3.26666667])
+catchment_min_1c = np.array([0.1, -4.5, -4.5, 4.6, 1.5])
+catchment_max_1c = np.array([5.5, 8.9, 3.3, 11.1, 5.1])

tests/_test_inputs/distance.py

@@ -100,3 +100,25 @@
         [-np.inf, 22.0, 6.0, 6.0, -np.inf],
     ]
 )
+
+# Distance to sink/source (no field weights, shortest path by number of edges)
+distance_1_to_sink_shortest = np.array(
+    [3.0, 3.0, 3.0, 3.0, 3.0, 2.0, 2.0, 2.0, 2.0, 3.0, 1.0, 1.0, 1.0, 2.0, 3.0, 1.0, 0.0,
+     1.0, 2.0, 3.0]
+)
+
+distance_1_to_source_shortest = np.array(
+    [0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 0.0, 2.0, 2.0, 2.0, 1.0,
+     0.0, 0.0, 1.0, 2.0, 1.0, 0.0]
+)
+
+# Length to sink (with field weights, shortest path)
+length_1_to_sink_shortest = np.array(
+    [22.0, 6.0, 10.0, 11.0, 12.0, 16.0, 5.0, 8.0, 8.0, 12.0, 9.0, 4.0, 3.0,
+     12.0, 21.0, 11.0, 3.0, 3.0, 9.0, 13.0]
+)
+
+length_1_to_source_shortest = np.array(
+    [6.0, 1.0, 2.0, 3.0, 4.0, 13.0, 2.0, 7.0, 9.0, 0.0, 19.0, 3.0, 7.0,
+     9.0, 9.0, 8.0, 6.0, 10.0, 10.0, 4.0]
+)

tests/catchments/array/test_max.py

@@ -0,0 +1,31 @@
+import numpy as np
+import pytest
+from _test_inputs.catchment import *
+from _test_inputs.accumulation import input_field_1c
+from _test_inputs.readers import *
+
+import earthkit.hydro as ekh
+
+
+@pytest.mark.parametrize(
+    "river_network, field, locations, expected",
+    [
+        (
+            ("cama_nextxy", cama_nextxy_1),
+            input_field_1c,
+            catchment_query_field_1,
+            catchment_max_1c,
+        ),
+    ],
+    indirect=["river_network"],
+)
+@pytest.mark.parametrize("array_backend", ["numpy", "torch"])
+def test_catchments_max(river_network, field, locations, expected, array_backend):
+    """Test catchment max aggregation."""
+    river_network = river_network.to_device("cpu", array_backend)
+    xp = ekh._backends.find.get_array_backend(array_backend)
+    result = ekh.catchments.array.max(river_network, xp.asarray(field), locations=locations)
+    result = np.asarray(result)
+    print("Result:", result)
+    print("Expected:", expected)
+    np.testing.assert_allclose(result, expected, rtol=1e-6)

tests/catchments/array/test_mean.py

@@ -0,0 +1,31 @@
+import numpy as np
+import pytest
+from _test_inputs.catchment import *
+from _test_inputs.accumulation import input_field_1c
+from _test_inputs.readers import *
+
+import earthkit.hydro as ekh
+
+
+@pytest.mark.parametrize(
+    "river_network, field, locations, expected",
+    [
+        (
+            ("cama_nextxy", cama_nextxy_1),
+            input_field_1c,
+            catchment_query_field_1,
+            catchment_mean_1c,
+        ),
+    ],
+    indirect=["river_network"],
+)
+@pytest.mark.parametrize("array_backend", ["numpy", "torch", "jax"])
+def test_catchments_mean(river_network, field, locations, expected, array_backend):
+    """Test catchment mean aggregation."""
+    river_network = river_network.to_device("cpu", array_backend)
+    xp = ekh._backends.find.get_array_backend(array_backend)
+    result = ekh.catchments.array.mean(river_network, xp.asarray(field), locations=locations)
+    result = np.asarray(result)
+    print("Result:", result)
+    print("Expected:", expected)
+    np.testing.assert_allclose(result, expected, rtol=1e-6)

tests/catchments/array/test_min.py

@@ -0,0 +1,31 @@
+import numpy as np
+import pytest
+from _test_inputs.catchment import *
+from _test_inputs.accumulation import input_field_1c
+from _test_inputs.readers import *
+
+import earthkit.hydro as ekh
+
+
+@pytest.mark.parametrize(
+    "river_network, field, locations, expected",
+    [
+        (
+            ("cama_nextxy", cama_nextxy_1),
+            input_field_1c,
+            catchment_query_field_1,
+            catchment_min_1c,
+        ),
+    ],
+    indirect=["river_network"],
+)
+@pytest.mark.parametrize("array_backend", ["numpy", "torch"])
+def test_catchments_min(river_network, field, locations, expected, array_backend):
+    """Test catchment min aggregation."""
+    river_network = river_network.to_device("cpu", array_backend)
+    xp = ekh._backends.find.get_array_backend(array_backend)
+    result = ekh.catchments.array.min(river_network, xp.asarray(field), locations=locations)
+    result = np.asarray(result)
+    print("Result:", result)
+    print("Expected:", expected)
+    np.testing.assert_allclose(result, expected, rtol=1e-6)

tests/catchments/array/test_sum.py

@@ -0,0 +1,31 @@
+import numpy as np
+import pytest
+from _test_inputs.catchment import *
+from _test_inputs.accumulation import input_field_1c
+from _test_inputs.readers import *
+
+import earthkit.hydro as ekh
+
+
+@pytest.mark.parametrize(
+    "river_network, field, locations, expected",
+    [
+        (
+            ("cama_nextxy", cama_nextxy_1),
+            input_field_1c,
+            catchment_query_field_1,
+            catchment_sum_1c,
+        ),
+    ],
+    indirect=["river_network"],
+)
+@pytest.mark.parametrize("array_backend", ["numpy", "torch", "jax"])
+def test_catchments_sum(river_network, field, locations, expected, array_backend):
+    """Test catchment sum aggregation."""
+    river_network = river_network.to_device("cpu", array_backend)
+    xp = ekh._backends.find.get_array_backend(array_backend)
+    result = ekh.catchments.array.sum(river_network, xp.asarray(field), locations=locations)
+    result = np.asarray(result)
+    print("Result:", result)
+    print("Expected:", expected)
+    np.testing.assert_allclose(result, expected, rtol=1e-6)

tests/distance/array/test_to_sink.py

@@ -0,0 +1,27 @@
+import numpy as np
+import pytest
+from _test_inputs.distance import *
+from _test_inputs.readers import *
+
+import earthkit.hydro as ekh
+
+
+@pytest.mark.parametrize(
+    "river_network, field, expected",
+    [
+        (
+            ("cama_nextxy", cama_nextxy_1),
+            None,
+            distance_1_to_sink_shortest,
+        ),
+    ],
+    indirect=["river_network"],
+)
+def test_distance_to_sink(river_network, field, expected):
+    """Test distance to sink computation."""
+    result = ekh.distance.array.to_sink(
+        river_network, field=field, path="shortest", return_type="masked"
+    )
+    print("Result:", result)
+    print("Expected:", expected)
+    np.testing.assert_array_equal(result, expected)

tests/distance/array/test_to_source.py

@@ -0,0 +1,27 @@
+import numpy as np
+import pytest
+from _test_inputs.distance import *
+from _test_inputs.readers import *
+
+import earthkit.hydro as ekh
+
+
+@pytest.mark.parametrize(
+    "river_network, field, expected",
+    [
+        (
+            ("cama_nextxy", cama_nextxy_1),
+            None,
+            distance_1_to_source_shortest,
+        ),
+    ],
+    indirect=["river_network"],
+)
+def test_distance_to_source(river_network, field, expected):
+    """Test distance to source computation."""
+    result = ekh.distance.array.to_source(
+        river_network, field=field, path="shortest", return_type="masked"
+    )
+    print("Result:", result)
+    print("Expected:", expected)
+    np.testing.assert_allclose(result, expected, rtol=1e-6)

tests/downstream/array/test_max.py

@@ -0,0 +1,39 @@
+import numpy as np
+import pytest
+from _test_inputs.accumulation import *
+from _test_inputs.readers import *
+
+import earthkit.hydro as ekh
+
+
+@pytest.mark.parametrize(
+    "river_network, input_field, flow_downstream, mv",
+    [
+        (
+            ("cama_nextxy", cama_nextxy_1),
+            input_field_1c,
+            downstream_metric_max_1c,
+            mv_1c,
+        ),
+        (
+            ("cama_nextxy", cama_nextxy_1),
+            input_field_1e,
+            downstream_metric_max_1e,
+            mv_1e,
+        ),
+    ],
+    indirect=["river_network"],
+)
+@pytest.mark.parametrize("array_backend", ["numpy", "torch", "jax"])
+def test_downstream_metric_max(river_network, input_field, flow_downstream, mv, array_backend):
+    river_network = river_network.to_device("cpu", array_backend)
+    xp = ekh._backends.find.get_array_backend(array_backend)
+    output_field = ekh.downstream.array.max(
+        river_network, xp.asarray(input_field), node_weights=None, return_type="masked"
+    )
+    output_field = np.asarray(output_field)
+    flow_downstream_out = np.asarray(xp.asarray(flow_downstream))
+    print(output_field)
+    print(flow_downstream_out)
+    assert output_field.dtype == flow_downstream_out.dtype
+    np.testing.assert_allclose(output_field, flow_downstream, rtol=1e-6, equal_nan=True)