diff --git a/tests/_test_inputs/accumulation.py b/tests/_test_inputs/accumulation.py index 9add82a6..7f6eb6cf 100644 --- a/tests/_test_inputs/accumulation.py +++ b/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, +) diff --git a/tests/_test_inputs/catchment.py b/tests/_test_inputs/catchment.py index 26e749f4..c9457b24 100644 --- a/tests/_test_inputs/catchment.py +++ b/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]) diff --git a/tests/_test_inputs/distance.py b/tests/_test_inputs/distance.py index f5f3e870..1d8bbf0c 100644 --- a/tests/_test_inputs/distance.py +++ b/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] +) diff --git a/tests/catchments/array/test_max.py b/tests/catchments/array/test_max.py new file mode 100644 index 00000000..ebd1686a --- /dev/null +++ b/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) diff --git a/tests/catchments/array/test_mean.py b/tests/catchments/array/test_mean.py new file mode 100644 index 00000000..817a7f56 --- /dev/null +++ b/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) diff --git a/tests/catchments/array/test_min.py b/tests/catchments/array/test_min.py new file mode 100644 index 00000000..8d0682b9 --- /dev/null +++ b/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) diff --git a/tests/catchments/array/test_sum.py b/tests/catchments/array/test_sum.py new file mode 100644 index 00000000..1b976128 --- /dev/null +++ b/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) diff --git a/tests/distance/array/test_to_sink.py b/tests/distance/array/test_to_sink.py new file mode 100644 index 00000000..0b11452c --- /dev/null +++ b/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) diff --git a/tests/distance/array/test_to_source.py b/tests/distance/array/test_to_source.py new file mode 100644 index 00000000..cd36ee89 --- /dev/null +++ b/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) diff --git a/tests/downstream/__init__.py b/tests/downstream/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/downstream/array/__init__.py b/tests/downstream/array/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/downstream/array/test_max.py b/tests/downstream/array/test_max.py new file mode 100644 index 00000000..2ca45d89 --- /dev/null +++ b/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) diff --git a/tests/downstream/array/test_mean.py b/tests/downstream/array/test_mean.py new file mode 100644 index 00000000..a2cb1a74 --- /dev/null +++ b/tests/downstream/array/test_mean.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_mean_1c, + mv_1c, + ), + ( + ("cama_nextxy", cama_nextxy_1), + input_field_1e, + downstream_metric_mean_1e, + mv_1e, + ), + ], + indirect=["river_network"], +) +@pytest.mark.parametrize("array_backend", ["numpy", "torch", "jax"]) +def test_downstream_metric_mean(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.mean( + 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) diff --git a/tests/downstream/array/test_min.py b/tests/downstream/array/test_min.py new file mode 100644 index 00000000..b5e6c92e --- /dev/null +++ b/tests/downstream/array/test_min.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_min_1c, + mv_1c, + ), + ( + ("cama_nextxy", cama_nextxy_1), + input_field_1e, + downstream_metric_min_1e, + mv_1e, + ), + ], + indirect=["river_network"], +) +@pytest.mark.parametrize("array_backend", ["numpy", "torch", "jax"]) +def test_downstream_metric_min(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.min( + 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) diff --git a/tests/downstream/array/test_sum.py b/tests/downstream/array/test_sum.py new file mode 100644 index 00000000..046a4d32 --- /dev/null +++ b/tests/downstream/array/test_sum.py @@ -0,0 +1,57 @@ +import numpy as np +import pytest +from _test_inputs.accumulation import * +from _test_inputs.readers import * +from utils import convert_to_2d + +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_sum_1c, + mv_1c, + ), + ( + ("cama_nextxy", cama_nextxy_1), + input_field_1e, + downstream_metric_sum_1e, + mv_1e, + ), + ], + indirect=["river_network"], +) +@pytest.mark.parametrize("array_backend", ["numpy", "torch", "jax"]) +def test_downstream_metric_sum( + 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.sum( + 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) + + print(input_field) + input_field = convert_to_2d(river_network, input_field, 0) + flow_downstream = convert_to_2d(river_network, flow_downstream, 0) + print(mv, input_field.dtype) + print(input_field, flow_downstream) + output_field = ekh.downstream.array.sum( + river_network, xp.asarray(input_field), node_weights=None + ) + output_field = np.asarray(output_field).flatten() + flow_downstream = np.asarray(xp.asarray(flow_downstream)) + print(output_field) + print(flow_downstream) + assert output_field.dtype == flow_downstream.dtype + np.testing.assert_allclose(output_field, flow_downstream, rtol=1e-6, equal_nan=True) diff --git a/tests/length/array/test_to_sink.py b/tests/length/array/test_to_sink.py new file mode 100644 index 00000000..abd10200 --- /dev/null +++ b/tests/length/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), + weights_1, + length_1_to_sink_shortest, + ), + ], + indirect=["river_network"], +) +def test_length_to_sink(river_network, field, expected): + """Test length to sink computation.""" + result = ekh.length.array.to_sink( + river_network, field=field, path="shortest", return_type="masked" + ) + print("Result:", result) + print("Expected:", expected) + np.testing.assert_allclose(result, expected, rtol=1e-6) diff --git a/tests/length/array/test_to_source.py b/tests/length/array/test_to_source.py new file mode 100644 index 00000000..9210268b --- /dev/null +++ b/tests/length/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), + weights_1, + length_1_to_source_shortest, + ), + ], + indirect=["river_network"], +) +def test_length_to_source(river_network, field, expected): + """Test length to source computation.""" + result = ekh.length.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) diff --git a/tests/subnetwork/__init__.py b/tests/subnetwork/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/subnetwork/test_subnetwork.py b/tests/subnetwork/test_subnetwork.py new file mode 100644 index 00000000..a0abfcd7 --- /dev/null +++ b/tests/subnetwork/test_subnetwork.py @@ -0,0 +1,111 @@ +import numpy as np +import pytest +from _test_inputs.readers import * + +import earthkit.hydro as ekh + + +@pytest.mark.parametrize( + "river_network", + [ + ("cama_nextxy", cama_nextxy_1), + ], + indirect=["river_network"], +) +def test_from_mask_node_mask(river_network): + """Test creating a subnetwork with a node mask.""" + # Create a node mask that selects a subset of nodes + node_mask = np.zeros(river_network.n_nodes, dtype=bool) + node_mask[:10] = True # Select first 10 nodes + + subnetwork = ekh.subnetwork.from_mask(river_network, node_mask=node_mask) + + # Check that the subnetwork has the correct number of nodes + assert subnetwork.n_nodes == 10 + assert subnetwork.n_nodes < river_network.n_nodes + + # Check that the subnetwork is a different object + assert subnetwork is not river_network + + +@pytest.mark.parametrize( + "river_network", + [ + ("cama_nextxy", cama_nextxy_1), + ], + indirect=["river_network"], +) +def test_from_mask_both_masks(river_network): + """Test creating a subnetwork with both node and edge masks.""" + # Create masks for both nodes and edges + node_mask = np.zeros(river_network.n_nodes, dtype=bool) + node_mask[:10] = True # Select first 10 nodes + + edge_mask = np.zeros(river_network.n_edges, dtype=bool) + edge_mask[:5] = True # Select first 5 edges + + subnetwork = ekh.subnetwork.from_mask( + river_network, node_mask=node_mask, edge_mask=edge_mask + ) + + # Check that the subnetwork has fewer nodes and edges + assert subnetwork.n_nodes <= 10 + assert subnetwork.n_edges <= 5 + assert subnetwork.n_nodes < river_network.n_nodes + + +@pytest.mark.parametrize( + "river_network", + [ + ("cama_nextxy", cama_nextxy_1), + ], + indirect=["river_network"], +) +def test_from_mask_no_mask(river_network): + """Test creating a subnetwork with no masks (should return a copy).""" + subnetwork = ekh.subnetwork.from_mask(river_network) + + # Check that it's a copy with the same properties + assert subnetwork.n_nodes == river_network.n_nodes + assert subnetwork.n_edges == river_network.n_edges + assert subnetwork is not river_network + + +@pytest.mark.parametrize( + "river_network", + [ + ("cama_nextxy", cama_nextxy_1), + ], + indirect=["river_network"], +) +def test_crop(river_network): + """Test cropping a gridded network to minimum bounding box.""" + # Skip test if river network doesn't have coords (required for crop) + if river_network._storage.coords is None: + pytest.skip("River network does not have coordinates required for crop") + + # First create a subnetwork with a specific node mask that leaves empty borders + # This ensures the crop will actually reduce the grid dimensions + node_mask = np.zeros(river_network.n_nodes, dtype=bool) + # Select nodes in the middle of the grid to leave empty rows/columns at edges + # Use a smaller subset to ensure there's space to crop + n_select = min(10, river_network.n_nodes // 2) + node_mask[:n_select] = True + + subnetwork = ekh.subnetwork.from_mask(river_network, node_mask=node_mask) + + # Skip if subnetwork doesn't have coords + if subnetwork._storage.coords is None: + pytest.skip("Subnetwork does not have coordinates required for crop") + + # Now crop the subnetwork + cropped = ekh.subnetwork.crop(subnetwork) + + # Check that number of nodes is preserved (crop doesn't remove nodes) + assert cropped.n_nodes == subnetwork.n_nodes + + # Check that the grid dimensions are reduced (actual cropping happened) + assert (cropped.shape[0] < subnetwork.shape[0]) or (cropped.shape[1] < subnetwork.shape[1]) + + # Check that it's a different object + assert cropped is not subnetwork diff --git a/tests/upstream/array/test_std.py b/tests/upstream/array/test_std.py new file mode 100644 index 00000000..e1c752ef --- /dev/null +++ b/tests/upstream/array/test_std.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, + upstream_metric_std_1c, + mv_1c, + ), + ( + ("cama_nextxy", cama_nextxy_1), + input_field_1e, + upstream_metric_std_1e, + mv_1e, + ), + ], + indirect=["river_network"], +) +@pytest.mark.parametrize("array_backend", ["numpy", "torch", "jax"]) +def test_calculate_upstream_metric_std(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.upstream.array.std( + 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-5, equal_nan=True) diff --git a/tests/upstream/array/test_var.py b/tests/upstream/array/test_var.py new file mode 100644 index 00000000..79a7a4d3 --- /dev/null +++ b/tests/upstream/array/test_var.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, + upstream_metric_var_1c, + mv_1c, + ), + ( + ("cama_nextxy", cama_nextxy_1), + input_field_1e, + upstream_metric_var_1e, + mv_1e, + ), + ], + indirect=["river_network"], +) +@pytest.mark.parametrize("array_backend", ["numpy", "torch", "jax"]) +def test_calculate_upstream_metric_var(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.upstream.array.var( + 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-5, equal_nan=True)