From 89e1f9967add1a99fffad68429f157dff728bc91 Mon Sep 17 00:00:00 2001
From: Bart de Koning <74617371+SouthEndMusic@users.noreply.github.com>
Date: Tue, 6 Aug 2024 10:52:56 +0200
Subject: [PATCH] Add `Makie` extension to Ribasim core (#1675)

Fixes https://github.com/Deltares/Ribasim/issues/1672.

I've added some basic functionality to plot storages and levels with
Makie:

![basin_data](https://github.com/user-attachments/assets/c86a9ea5-4d81-4df9-84a5-2b2c2f087901)


![flow](https://github.com/user-attachments/assets/dcbd9d08-deab-450c-847c-e37fff39c536)
(For conservative nodes I only plot the inflow)

I'm first going to work on
https://github.com/Deltares/Ribasim/issues/1674 and then visualize that.

---------

Co-authored-by: Martijn Visser <mgvisser@gmail.com>
---
 Manifest.toml               |  9 +++++-
 core/Project.toml           |  8 ++++++
 core/ext/RibasimMakieExt.jl | 56 +++++++++++++++++++++++++++++++++++++
 core/src/Ribasim.jl         |  6 ++++
 4 files changed, 78 insertions(+), 1 deletion(-)
 create mode 100644 core/ext/RibasimMakieExt.jl

diff --git a/Manifest.toml b/Manifest.toml
index 34a68402a..d7b27cb04 100644
--- a/Manifest.toml
+++ b/Manifest.toml
@@ -2,7 +2,7 @@
 
 julia_version = "1.10.4"
 manifest_format = "2.0"
-project_hash = "1f635113c3942011817a95f3b34347acc94b7036"
+project_hash = "4ab334c22b392689c0b97d7d9437877b605ab2cb"
 
 [[deps.ADTypes]]
 git-tree-sha1 = "aa4d425271a914d8c4af6ad9fccb6eb3aec662c7"
@@ -1302,6 +1302,13 @@ path = "core"
 uuid = "aac5e3d9-0b8f-4d4f-8241-b1a7a9632635"
 version = "2024.10.0"
 
+    [deps.Ribasim.extensions]
+    RibasimMakieExt = ["Makie", "DataFrames"]
+
+    [deps.Ribasim.weakdeps]
+    DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
+    Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
+
 [[deps.RuntimeGeneratedFunctions]]
 deps = ["ExprTools", "SHA", "Serialization"]
 git-tree-sha1 = "04c968137612c4a5629fa531334bb81ad5680f00"
diff --git a/core/Project.toml b/core/Project.toml
index 302293a24..1f1c27f1e 100644
--- a/core/Project.toml
+++ b/core/Project.toml
@@ -40,6 +40,13 @@ TerminalLoggers = "5d786b92-1e48-4d6f-9151-6b4477ca9bed"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 TranscodingStreams = "3bb67fe8-82b1-5028-8e26-92a6c54297fa"
 
+[weakdeps]
+Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
+DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
+
+[extensions]
+RibasimMakieExt = ["Makie", "DataFrames"]
+
 [compat]
 Accessors = "0.1"
 Aqua = "0.8"
@@ -67,6 +74,7 @@ Legolas = "0.5"
 LinearSolve = "2.24"
 Logging = "<0.0.1, 1"
 LoggingExtras = "1"
+Makie = "0.21"
 MetaGraphsNext = "0.6, 0.7"
 OrdinaryDiffEq = "6.7"
 PreallocationTools = "0.4"
diff --git a/core/ext/RibasimMakieExt.jl b/core/ext/RibasimMakieExt.jl
new file mode 100644
index 000000000..97077bab8
--- /dev/null
+++ b/core/ext/RibasimMakieExt.jl
@@ -0,0 +1,56 @@
+module RibasimMakieExt
+using DataFrames: DataFrame
+using Makie: Figure, Axis, lines!, axislegend
+using Ribasim: Ribasim, Model
+
+function Ribasim.plot_basin_data!(model::Model, ax::Axis, column::Symbol)
+    basin_data = DataFrame(Ribasim.basin_table(model))
+    for node_id in unique(basin_data.node_id)
+        group = filter(:node_id => ==(node_id), basin_data)
+        lines!(ax, group.time, getproperty(group, column); label = "Basin #$node_id")
+    end
+
+    axislegend(ax)
+    return nothing
+end
+
+function Ribasim.plot_basin_data(model::Model)
+    f = Figure()
+    ax1 = Axis(f[1, 1]; ylabel = "level [m]")
+    ax2 = Axis(f[2, 1]; xlabel = "time", ylabel = "storage [m³]")
+    Ribasim.plot_basin_data!(model, ax1, :level)
+    Ribasim.plot_basin_data!(model, ax2, :storage)
+    f
+end
+
+function Ribasim.plot_flow!(
+    model::Model,
+    ax::Axis,
+    edge_id::Int32;
+    skip_conservative_out = false,
+)
+    flow_data = DataFrame(Ribasim.flow_table(model))
+    flow_data = filter(:edge_id => ==(edge_id), flow_data)
+    first_row = first(flow_data)
+    # Skip outflows of conservative nodes because these are the same as the inflows
+    if skip_conservative_out &&
+       Ribasim.NodeType.T(first_row.from_node_type) in Ribasim.conservative_nodetypes
+        return nothing
+    end
+    label = "$(first_row.from_node_type) #$(first_row.from_node_id) → $(first_row.to_node_type) #$(first_row.to_node_id)"
+    lines!(ax, flow_data.time, flow_data.flow_rate; label)
+    return nothing
+end
+
+function Ribasim.plot_flow(model::Model)
+    f = Figure()
+    ax = Axis(f[1, 1]; xlabel = "time", ylabel = "flow rate [m³s⁻¹]")
+    edge_ids = unique(Ribasim.flow_table(model).edge_id)
+    for edge_id in edge_ids
+        Ribasim.plot_flow!(model, ax, edge_id; skip_conservative_out = true)
+    end
+    axislegend(ax)
+    f
+end
+
+end # module RibasimMakieExt
diff --git a/core/src/Ribasim.jl b/core/src/Ribasim.jl
index f64aaa808..459a1ae73 100644
--- a/core/src/Ribasim.jl
+++ b/core/src/Ribasim.jl
@@ -98,4 +98,10 @@ include("callback.jl")
 include("main.jl")
 include("libribasim.jl")
 
+# Define names used in Makie extension
+function plot_basin_data end
+function plot_basin_data! end
+function plot_flow end
+function plot_flow! end
+
 end  # module Ribasim