Merge pull request #26 from ecmwf/prep_for_v1

Prep for v1

Author: Oisin-M

README.md

@@ -49,6 +49,11 @@ import earthkit.hydro as ekh
 The package contains different ways of constructing or loading a `RiverNetwork` object. A `RiverNetwork` object is a representation of a river network on a grid.
 It can be used to compute basic hydrological functions, such as propagating a scalar along the river network or extract a catchment from the river network.
 
+### Mathematical Details
+Given a discretisation of a domain i.e. a set of points $\mathcal{D}=\{ (x_i, y_i)\}_{i=1}^N$, a river network is a directed acyclic graph $\mathcal{R}=(V,E)$ where the vertices $V \subseteq \mathcal{D}$. The out-degree of each vertex is at most 1 i.e. each point in the river network points to at most one downstream location.
+
+For ease of notation, if an edge exists from $(x_i, y_i)$ to $(x_j, y_j)$, we write $i \rightarrow j$.
+
 ### Readers
 
 ```
@@ -79,31 +84,36 @@ Creates a `RiverNetwork` from a CaMa-Flood bin format of type "downxy" or "nextx
 ```
 network.accuflux(field)
 ```
-Calculates the total accumulated flux down a river network.
+Calculates the total accumulated flux down a river network.\
+$$v_i^{\prime}=v_i+\sum_{j \rightarrow i}~v_j^{\prime}$$
 
 <img src="docs/images/accuflux.gif" width="200px" height="160px" />
 
 ```
 network.upstream(field)
 ```
-Updates each node with the sum of its upstream nodes.
+Updates each node with the sum of its upstream nodes.\
+$$v_i^{\prime}=\sum_{j \rightarrow i}~v_j$$
 
 ```
 network.downstream(field)
 ```
-Updates each node with its downstream node.
+Updates each node with its downstream node.\
+$$v_i^{\prime} = v_j, ~j ~ \text{s.t.} ~ i \rightarrow j$$
 
 ```
 network.catchment(field)
 ```
-Finds the catchments (all upstream nodes of specified nodes, with overwriting).
+Finds the catchments (all upstream nodes of specified nodes, with overwriting).\
+$$v_i^{\prime} = v_j^{\prime}  ~ \text{if} ~  v_j^{\prime} \neq 0 ~ \text{else} ~ v_i, ~j ~ \text{s.t.} ~ i \rightarrow j$$
 
 <img src="docs/images/catchment.gif" width="200px" height="160px" />
 
 ```
 network.subcatchment(field)
 ```
-Finds the subcatchments (all upstream nodes of specified nodes, without overwriting).
+Finds the subcatchments (all upstream nodes of specified nodes, without overwriting).\
+$$v_i^{\prime} = v_j^{\prime}  ~ \text{if} ~  (v_j^{\prime} \neq 0 ~ \text{and} ~ v_j = 0) ~ \text{else} ~ v_i, ~j ~ \text{s.t.} ~ i \rightarrow j$$
 
 <img src="docs/images/subcatchment.gif" width="200px" height="160px" />
 

pyproject.toml

@@ -29,7 +29,6 @@ dynamic = ["version", "readme"]
 
 dependencies = [
     "numpy",
-    "xarray",
     "joblib"
 ]
 
@@ -39,6 +38,9 @@ dependencies = [
     issues = "https://github.com/ecmwf/earthkit-hydro/issues"
 
 [project.optional-dependencies]
+    netcdf = [
+        "earthkit-data",
+    ]
     test = [
         "pytest",
         "pytest-cases",

src/earthkit/hydro/readers.py

@@ -1,5 +1,4 @@
 import numpy as np
-import xarray as xr
 from .river_network import RiverNetwork
 import joblib
 from .caching import Cache
@@ -14,7 +13,7 @@ def load_river_network(
     cache_fname="{ekh_version}_{domain}_{version}.joblib",
 ):
     """
-    Loads a river network from a specified.
+    Loads a precomputed river network.
     A cache is used to store the river network file locally.
 
     Parameters
@@ -44,6 +43,24 @@ def load_river_network(
     return network
 
 
+def load_saved_network(filepath):
+    """
+    Loads a saved river network from file.
+
+    Parameters
+    ----------
+    filepath : str
+        Path to the saved river network joblib file.
+
+    Returns
+    -------
+    RiverNetwork
+        The loaded river network object.
+    """
+    network = joblib.load(filepath)
+    return network
+
+
 def from_netcdf_d8(filename):
     """
     Loads a river network from a NetCDF file using D8 direction encoding.
@@ -58,7 +75,13 @@ def from_netcdf_d8(filename):
     RiverNetwork
         The constructed river network object.
     """
-    data = xr.open_dataset(filename, mask_and_scale=False)["Band1"].values
+    try:
+        import earthkit.data as ekd
+    except ModuleNotFoundError:
+        raise ModuleNotFoundError(
+            "earthkit-data is required for netcdf support.\nTo install it, run `pip install earthkit-data`"
+        )
+    data = ekd.from_source("file", filename).to_xarray(mask_and_scale=False)["Band1"].values
     return from_d8(data)
 
 
@@ -96,7 +119,7 @@ def from_d8(data):
     return create_network(upstream_indices, downstream_indices, missing_mask, shape)
 
 
-def from_netcdf_cama(filename, type="downxy"):
+def from_netcdf_cama(filename, type="nextxy"):
     """
     Loads a river network from a CaMa-Flood style NetCDF file.
 
@@ -117,10 +140,16 @@ def from_netcdf_cama(filename, type="downxy"):
     Exception
         If an unknown type is specified.
     """
-    data = xr.open_dataset(filename, mask_and_scale=False)
+    try:
+        import earthkit.data as ekd
+    except ModuleNotFoundError:
+        raise ModuleNotFoundError(
+            "earthkit-data is required for netcdf support.\nTo install it, run `pip install earthkit-data`"
+        )
+    data = ekd.from_source("file", filename).to_xarray(mask_and_scale=False)
     if type == "downxy":
         dx, dy = data.downx.values, data.downy.values
-        return from_cama_downxy(dx, -dy)
+        return from_cama_downxy(dx, dy)
     elif type == "nextxy":
         x, y = data.nextx.values, data.nexty.values
         return from_cama_nextxy(x, y)
@@ -199,11 +228,11 @@ def from_cama_downxy(dx, dy):
     x_offsets = x_offsets[mask_upstream]
     y_offsets = y_offsets[mask_upstream]
 
-    upstream_indicies, downstream_indices = find_upstream_downstream_indices_from_offsets(
+    upstream_indices, downstream_indices = find_upstream_downstream_indices_from_offsets(
         x_offsets, y_offsets, missing_mask, mask_upstream, shape
     )
 
-    return create_network(upstream_indicies, downstream_indices, missing_mask, shape)
+    return create_network(upstream_indices, downstream_indices, missing_mask, shape)
 
 
 def from_cama_nextxy(x, y):

src/earthkit/hydro/river_network.py

@@ -191,7 +191,9 @@ class RiverNetwork:
         Groups of nodes sorted in topological order.
     """
 
-    def __init__(self, nodes, downstream, mask, sinks=None, sources=None, topological_labels=None) -> None:
+    def __init__(
+        self, nodes, downstream, mask, sinks=None, sources=None, topological_labels=None, check_for_cycles=False
+    ) -> None:
         """
         Initialises the RiverNetwork with nodes, downstream nodes, and a mask.
 
@@ -203,6 +205,14 @@ def __init__(self, nodes, downstream, mask, sinks=None, sources=None, topologica
             Array of downstream node ids corresponding to each node.
         mask : numpy.ndarray
             A mask converting from the domain to the river network.
+        sinks : numpy.ndarray, optional
+            Array of sinks of the river network.
+        sources :  numpy.ndarray, optional
+            Array of sources of the river network.
+        topological_labels : numpy.ndarray, optional
+            Array of precomputed topological distance labels.
+        check_for_cycles : bool, optional
+            Whether to check for cycles when instantiating the river network.
         """
         self.nodes = nodes
         self.n_nodes = len(nodes)
@@ -212,6 +222,8 @@ def __init__(self, nodes, downstream, mask, sinks=None, sources=None, topologica
             sinks if sinks is not None else self.nodes[self.downstream_nodes == self.n_nodes]
         )  # nodes with no downstreams
         self.sources = sources if sources is not None else self.get_sources()  # nodes with no upstreams
+        if check_for_cycles:
+            self.check_for_cycles()
         self.topological_labels = (
             topological_labels if topological_labels is not None else self.compute_topological_labels()
         )
@@ -275,6 +287,19 @@ def get_sources(self):
         inlets = tmp_nodes[tmp_nodes != -1]  # sources are nodes that are not downstream nodes
         return inlets
 
+    def check_for_cycles(self):
+        """
+        Checks if the river network contains any cycles and raises an Exception if it does.
+        """
+        nodes = self.downstream_nodes.copy()
+        while True:
+            if np.any(nodes == self.nodes):
+                Exception("River Network contains a cycle.")
+            elif np.all(nodes == self.n_nodes):
+                break
+            not_sinks = nodes != self.n_nodes
+            nodes[not_sinks] = self.downstream_nodes[nodes[not_sinks]].copy()
+
     def compute_topological_labels(self):
         """
         Finds the topological distance labels for each node in the river network.
@@ -290,6 +315,8 @@ def compute_topological_labels(self):
         current_sum = 0  # sum of labels
         n = 1  # distance from source
         while current_sum > old_sum:
+            if n > self.n_nodes:
+                raise Exception("River Network contains a cycle.")
             old_sum = current_sum
             inlets = inlets[inlets != self.n_nodes]  # subset to valid nodes
             labels[inlets] = n  # update furthest distance from source