Feat/add plotter

README.md

@@ -105,30 +105,27 @@ Creates a `RiverNetwork`. Current options are
 - source: An earthkit-data compatable source. See [list](https://earthkit-data.readthedocs.io/en/latest/guide/sources.html)
 
 ### Computing Metrics Over River Networks
+_Currently supported metrics are "sum", "mean", "max", "min" and "product". If weights is provided, it is used to weight the field in the calculation._
 
 ```
 ekh.calculate_catchment_metric(river_network, field, stations, metric, weights=None)
 ```
-Calculates the metric over each catchment defined by stations. Current options are
-- metric: "sum", "max", "min", "mean"
+Calculates the metric over each catchment defined by stations.
 
 ```
 ekh.calculate_subcatchment_metric(river_network, field, stations, metric, weights=None)
 ```
-Calculates the metric over each subcatchment defined by stations. Current options are
-- metric: "sum", "max", "min", "mean"
+Calculates the metric over each subcatchment defined by stations.
 
 ```
 ekh.calculate_upstream_metric(river_network, field, metric, weights=None)
 ```
-Calculates a metric over all upstream nodes for a river network. If weights is provided, it is used to weight the field in the calculation. Options are
-- metric: "sum", "mean", "max", "min"
+Calculates a metric over all upstream nodes for a river network.
 
 ```
 ekh.calculate_metric_for_labels(field, labels, metric, weights=None)
 ```
-_(for advanced users)_ Calculates a metric over field for each label in the labels field. If weights is provided, it is used to weight the field in the calculation. Options are
-- metric: "sum", "mean", "max", "min"
+_(for advanced users)_ Calculates a metric over field for each label in the labels field.
 
 ```
 ekh.flow_downstream(river_network, field)

src/earthkit/hydro/__init__.py

@@ -4,6 +4,7 @@
 from .core import flow
 from .label import calculate_metric_for_labels
 from .movement import move_downstream, move_upstream
+from .plot import plot
 from .readers import (
     create_river_network,
     from_cama_downxy,

src/earthkit/hydro/metrics.py

@@ -21,4 +21,9 @@ class Min:
     base_val = np.inf
 
 
-metrics_dict = {"sum": Sum, "mean": Mean, "max": Max, "min": Min}
+class Product:
+    func = np.multiply
+    base_val = 1
+
+
+metrics_dict = {"sum": Sum, "mean": Mean, "max": Max, "min": Min, "product": Product}

src/earthkit/hydro/plot.py

@@ -0,0 +1,33 @@
+import numpy as np
+
+
+def plot(river_network, field, ax=None):
+
+    try:
+        import matplotlib.pyplot as plt
+    except ModuleNotFoundError:
+        raise ModuleNotFoundError(
+            "matplotlib.pyplot is required for plotting."
+            "\nTo install it, run `pip install matplotlib`"
+        )
+
+    x = np.arange(river_network.mask.shape[1])
+    y = np.arange(river_network.mask.shape[0])[::-1]
+
+    X, Y = np.meshgrid(x, y)
+    x = X[river_network.mask]
+    y = Y[river_network.mask]
+
+    not_sinks = river_network.downstream_nodes != river_network.n_nodes
+    u = np.zeros_like(x)
+    u[not_sinks] = x[river_network.downstream_nodes[not_sinks]] - x[not_sinks]
+    v = np.zeros_like(x)
+    v[not_sinks] = y[river_network.downstream_nodes[not_sinks]] - y[not_sinks]
+
+    c = field[river_network.mask]
+
+    if ax is None:
+        _, ax = plt.subplots()
+
+    ax.quiver(x, y, u, v, c, scale=1, scale_units="xy", angles="xy")
+    return ax