Document parsing Delwaq results.

docs/guide/delwaq.ipynb

@@ -54,6 +54,29 @@
     "model.plot();  # for later comparison"
    ]
   },
+  {
+   "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 +112,7 @@
    "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",
     "graph, substances = generate(toml_path, output_path)"
    ]
   },
@@ -182,6 +203,70 @@
     "- 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)"
+   ]
+  },
+  {
+   "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, 9, [\"Foo\", \"Bar\"])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from ribasim.delwaq import plot_spatial\n",
+    "\n",
+    "plot_spatial(nmodel, \"FlowBoundary\")"
+   ]
   }
  ],
  "metadata": {
@@ -200,7 +285,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 @@ def generate(
     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,12 @@ def generate(
     # 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,
     }
     substances.update(defaults.keys())
 
@@ -486,6 +490,33 @@ def generate(
     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,57 @@
-def plot():
-    pass
+import matplotlib.pyplot as plt
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+
+
+def plot_fraction(
+    model,
+    node_id,
+    tracers=["Basin", "LevelBoundary", "FlowBoundary", "UserDemand", "Initial"],
+):
+    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.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"):
+    table = model.basin.concentration_external.df
+    table = table[table["substance"] == tracer]
+    table = table[table["time"] == table["time"].max()]
+    table.set_index("node_id", inplace=True)
+
+    nodes = model.node_table().df
+    nodes = nodes[nodes.index.isin(table.index)]
+
+    fig, ax = plt.subplots()
+    s = ax.scatter(
+        nodes.geometry.x,
+        nodes.geometry.y,
+        c=table["concentration"][nodes.index],
+        clim=(0, 1),
+    )
+    ax.legend()
+    ax.set_title(f"Scatter plot for {tracer} tracer at {table["time"].iloc[0]}")
+
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes("right", size="5%", pad=0.05)
+
+    fig.colorbar(s, cax=cax, orientation="vertical")
+    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)
@@ -27,6 +32,7 @@ def test_offline_delwaq_coupling():
         "Cl",
         "Continuity",
         "FlowBoundary",
+        "Foo",
         "LevelBoundary",
         "Terminal",
         "Tracer",