From a310f1982f45ed1307df7e2e5f99a0b6ca421600 Mon Sep 17 00:00:00 2001 From: Maarten Pronk Date: Tue, 24 Sep 2024 08:09:49 +0200 Subject: [PATCH] Document parsing Delwaq results. --- docs/guide/delwaq.ipynb | 93 ++++++++++++++++++++++- python/ribasim/ribasim/delwaq/__init__.py | 14 +++- python/ribasim/ribasim/delwaq/generate.py | 37 ++++++++- python/ribasim/ribasim/delwaq/plot.py | 59 +++++++++++++- python/ribasim/tests/test_delwaq.py | 8 +- 5 files changed, 198 insertions(+), 13 deletions(-) diff --git a/docs/guide/delwaq.ipynb b/docs/guide/delwaq.ipynb index 1082aba42..d72519ed5 100644 --- a/docs/guide/delwaq.ipynb +++ b/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, diff --git a/python/ribasim/ribasim/delwaq/__init__.py b/python/ribasim/ribasim/delwaq/__init__.py index 9a7ac18f0..9aaec6658 100644 --- a/python/ribasim/ribasim/delwaq/__init__.py +++ b/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", +] diff --git a/python/ribasim/ribasim/delwaq/generate.py b/python/ribasim/ribasim/delwaq/generate.py index 740d5cd42..94453625e 100644 --- a/python/ribasim/ribasim/delwaq/generate.py +++ b/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] diff --git a/python/ribasim/ribasim/delwaq/plot.py b/python/ribasim/ribasim/delwaq/plot.py index 52a508221..70406c00a 100644 --- a/python/ribasim/ribasim/delwaq/plot.py +++ b/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) diff --git a/python/ribasim/tests/test_delwaq.py b/python/ribasim/tests/test_delwaq.py index bedbcb4e4..209e5a716 100644 --- a/python/ribasim/tests/test_delwaq.py +++ b/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",