Document parsing Delwaq results.

docs/guide/delwaq.ipynb

@@ -27,6 +27,7 @@
     "from pathlib import Path\n",
     "\n",
     "toml_path = Path(\"../../generated_testmodels/basic/ribasim.toml\")\n",
+    "\n",
     "assert toml_path.is_file()"
    ]
   },
@@ -54,6 +55,38 @@
     "model.plot();  # for later comparison"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.basin.profile"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Let's add another tracer to the model, to setup a fraction calculation."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from ribasim.delwaq import add_tracer\n",
+    "\n",
+    "add_tracer(model, 11, \"Foo\")\n",
+    "add_tracer(model, 15, \"Bar\")\n",
+    "display(model.flow_boundary.concentration)  # flow boundaries\n",
+    "display(model.level_boundary.concentration)  # flow boundaries\n",
+    "\n",
+    "model.write(toml_path)"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -89,9 +122,8 @@
    "source": [
     "from ribasim.delwaq import generate\n",
     "\n",
-    "output_path = Path(\n",
-    "    \"../../generated_testmodels/basic/delwaq\"\n",
-    ")  # set a path where we store the Delwaq input files\n",
+    "output_path = Path(\"../../generated_testmodels/basic/delwaq\")\n",
+    "\n",
     "graph, substances = generate(toml_path, output_path)"
    ]
   },
@@ -182,6 +214,83 @@
     "- Basin boundaries are split into separate nodes and links (drainage, precipitation, and evaporation, as indicated by the duplicated Basin IDs on the right hand side)\n",
     "- All node IDs have been renumbered, with boundaries being negative, and Basins being positive."
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Parsing the results\n",
+    "With Delwaq having run, we can now parse the results using `ribasim.delwaq.parse`. This function requires the `graph` and `substances` variables that were output by `ribasim.delwaq.generate`, as well as the path to the results folder of the Delwaq simulation."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from ribasim.delwaq import parse\n",
+    "\n",
+    "nmodel = parse(toml_path, graph, substances, output_folder=output_path)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The parsed model is identical to the Ribasim model, with the exception of the added concentration_external table that contains all tracer results from Delwaq."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "display(nmodel.basin.concentration_external)\n",
+    "print(substances)\n",
+    "t = nmodel.basin.concentration_external.df\n",
+    "t[t.time == t.time.unique()[2]]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We can use this table to plot the results of the Delwaq model, both spatially as over time."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from ribasim.delwaq import plot_fraction\n",
+    "\n",
+    "plot_fraction(nmodel, 1)  # default tracers, should add up to 1\n",
+    "plot_fraction(nmodel, 9, [\"Foo\", \"Bar\"])  # custom tracers\n",
+    "plot_fraction(nmodel, 9, [\"Continuity\"])  # mass balance check"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from ribasim.delwaq import plot_spatial\n",
+    "\n",
+    "plot_spatial(nmodel, \"Bar\")\n",
+    "plot_spatial(nmodel, \"Foo\", versus=\"Bar\")  # ratio of Meuse to Rhine"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {
@@ -200,7 +309,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.4"
+   "version": "3.12.5"
   }
  },
  "nbformat": 4,

python/ribasim/ribasim/delwaq/__init__.py

@@ -1,6 +1,14 @@
-from .generate import generate
+from .generate import add_tracer, generate
 from .parse import parse
-from .plot import plot
+from .plot import plot_fraction, plot_spatial
 from .util import run_delwaq
 
-__all__ = ["generate", "parse", "run_delwaq", "plot"]
+__all__ = [
+    "generate",
+    "parse",
+    "run_delwaq",
+    "plot",
+    "add_tracer",
+    "plot_fraction",
+    "plot_spatial",
+]

python/ribasim/ribasim/delwaq/generate.py

@@ -5,7 +5,8 @@
 from datetime import timedelta
 from pathlib import Path
 
-from ribasim.utils import MissingOptionalModule
+from ribasim import nodes
+from ribasim.utils import MissingOptionalModule, _pascal_to_snake
 
 try:
     import networkx as nx
@@ -276,13 +277,14 @@
     toml_path: Path,
     output_folder=output_folder,
     use_evaporation=USE_EVAP,
+    results_folder="results",
 ) -> tuple[nx.DiGraph, set[str]]:
     """Generate a Delwaq model from a Ribasim model and results."""
 
     # Read in model and results
     model = ribasim.Model.read(toml_path)
-    basins = pd.read_feather(toml_path.parent / "results" / "basin.arrow")
-    flows = pd.read_feather(toml_path.parent / "results" / "flow.arrow")
+    basins = pd.read_feather(toml_path.parent / results_folder / "basin.arrow")
+    flows = pd.read_feather(toml_path.parent / results_folder / "flow.arrow")
 
     output_folder.mkdir(exist_ok=True)
 
@@ -409,10 +411,14 @@
     # Setup initial basin concentrations
     defaults = {
         "Continuity": 1.0,
+        "Initial": 1.0,
         "Basin": 0.0,
         "LevelBoundary": 0.0,
         "FlowBoundary": 0.0,
         "Terminal": 0.0,
+        "UserDemand": 0.0,
+        "Precipitation": 0.0,
+        "Drainage": 0.0,
     }
     substances.update(defaults.keys())
 
@@ -486,6 +492,33 @@
     return G, substances
 
 
+def add_tracer(model, node_id, tracer_name):
+    """Add a tracer to the Delwaq model."""
+    n = model.node_table().df.loc[node_id]
+    node_type = n.node_type
+    if node_type not in [
+        "Basin",
+        "LevelBoundary",
+        "FlowBoundary",
+        "UserDemand",
+    ]:
+        raise ValueError("Can only trace Basins and boundaries")
+    snake_node_type = _pascal_to_snake(node_type)
+    nt = getattr(model, snake_node_type)
+
+    ct = getattr(nodes, snake_node_type)
+    table = ct.Concentration(
+        node_id=[node_id],
+        time=[model.starttime],
+        substance=[tracer_name],
+        concentration=[1.0],
+    )
+    if nt.concentration is None:
+        nt.concentration = table
+    else:
+        nt.concentration = pd.concat([nt.concentration.df, table.df], ignore_index=True)
+
+
 if __name__ == "__main__":
     # Generate a Delwaq model from the default Ribasim model
     repo_dir = delwaq_dir.parents[1]

python/ribasim/ribasim/delwaq/plot.py

@@ -1,2 +1,94 @@
-def plot():
-    pass
+import matplotlib.pyplot as plt
+import numpy as np
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+
+
+def plot_fraction(
+    model,
+    node_id,
+    tracers=[
+        "Basin",
+        "LevelBoundary",
+        "FlowBoundary",
+        "UserDemand",
+        "Initial",
+        "Drainage",
+        "Precipitation",
+        "Terminal",
+    ],
+):
+    table = model.basin.concentration_external.df
+    table = table[table["node_id"] == node_id]
+    table = table[table["substance"].isin(tracers)]
+
+    groups = table.groupby("substance")
+    stack = {k: v["concentration"].to_numpy() for (k, v) in groups}
+
+    fig, ax = plt.subplots()
+    ax.stackplot(
+        groups.get_group(tracers[0])["time"],
+        stack.values(),
+        labels=stack.keys(),
+    )
+    ax.plot(
+        groups.get_group(tracers[0])["time"],
+        np.sum(list(stack.values()), axis=0),
+        c="black",
+        lw=2,
+    )
+    ax.legend()
+    ax.set_title(f"Fraction plot for node {node_id}")
+    ax.set_xlabel("Time")
+    ax.set_ylabel("Fraction")
+
+    plt.show(fig)
+
+
+def plot_spatial(model, tracer="Basin", versus=None, limit=0.001):
+    table = model.basin.concentration_external.df
+    table = table[table["time"] == table["time"].max()]
+
+    if versus is not None:
+        vtable = table[table["substance"] == versus]
+        vtable.set_index("node_id", inplace=True)
+    table = table[table["substance"] == tracer]
+    table.set_index("node_id", inplace=True)
+
+    nodes = model.node_table().df
+    nodes = nodes[nodes.index.isin(table.index)]
+
+    if versus is None:
+        c = table["concentration"][nodes.index]
+        alpha = c > limit
+    else:
+        alpha = (
+            table["concentration"][nodes.index] + vtable["concentration"][nodes.index]
+        )
+        c = table["concentration"][nodes.index] / alpha
+
+    fig, ax = plt.subplots()
+    s = ax.scatter(
+        nodes.geometry.x,
+        nodes.geometry.y,
+        c=c,
+        clim=(0, 1),
+        alpha=alpha,
+    )
+    dt = table["time"].iloc[0]
+    if versus is None:
+        ax.set_title(f"Scatter plot for {tracer} tracer at {dt}")
+    else:
+        ax.set_title(f"Scatter plot for {tracer} vs {versus} tracer at {dt}")
+
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes("right", size="5%", pad=0.05)
+
+    fig.colorbar(s, cax=cax, orientation="vertical")
+    if versus is not None:
+        cax.set_ylabel(f"{tracer} fraction vs {versus} fraction")
+    else:
+        cax.set_ylabel(f"Overall {tracer} fraction")
+    ax.set_xlabel("x")
+    ax.set_ylabel("y")
+
+    plt.show(fig)

python/ribasim/tests/test_delwaq.py

@@ -2,7 +2,8 @@
 from pathlib import Path
 
 import pytest
-from ribasim.delwaq import generate, parse, run_delwaq
+from ribasim import Model
+from ribasim.delwaq import add_tracer, generate, parse, run_delwaq
 
 delwaq_dir = Path(__file__).parent
 
@@ -14,6 +15,10 @@ def test_offline_delwaq_coupling():
     repo_dir = delwaq_dir.parents[2]
     toml_path = repo_dir / "generated_testmodels/basic/ribasim.toml"
 
+    model = Model.read(toml_path)
+    add_tracer(model, 17, "Foo")
+    model.write(toml_path)
+
     graph, substances = generate(toml_path)
     run_delwaq()
     model = parse(toml_path, graph, substances)
@@ -26,8 +31,13 @@ def test_offline_delwaq_coupling():
         "Basin",
         "Cl",
         "Continuity",
+        "Drainage",
         "FlowBoundary",
+        "Foo",
+        "Initial",
         "LevelBoundary",
+        "Precipitation",
         "Terminal",
         "Tracer",
+        "UserDemand",
     ]