From 92c6f11f5b661641715ae11b7065aba8b1a223b4 Mon Sep 17 00:00:00 2001 From: Bart de Koning Date: Thu, 31 Oct 2024 07:59:22 +0100 Subject: [PATCH] Code reorganisation and water_balance! example --- docs/dev/callstacks.qmd | 203 +++++++++++++++++++++++----------------- 1 file changed, 115 insertions(+), 88 deletions(-) diff --git a/docs/dev/callstacks.qmd b/docs/dev/callstacks.qmd index 4bfd6b8ff..ec6d02068 100644 --- a/docs/dev/callstacks.qmd +++ b/docs/dev/callstacks.qmd @@ -106,51 +106,31 @@ end ############################################# -using Ribasim -using SQLite -``` - -## Parameter initialization - -```{julia} -toml_path = normpath(@__DIR__, "../../generated_testmodels/basic_transient/ribasim.toml") -config = Ribasim.Config(toml_path) -db_path = Ribasim.database_path(config) -db = SQLite.DB(db_path) - -graph_orig, verts = tracecall((Ribasim,), Ribasim.Parameters, (db, config)) -close(db) -``` - -```{julia} -graph = copy(graph_orig) - using GLMakie using Colors -f = Figure() -ax = Axis(f[1, 1]) +function get_node_depths(graph) + depths = dijkstra_shortest_paths(graph, 1).dists + nodes_per_depth = Dict(Int(depth) => Int[] for depth in unique(depths)) -# Options -maxdepth = 5 -plot_non_Ribasim = false -squash_per_depth = true + for (i, depth) in enumerate(depths) + nm = graph[i] + nm.depth[] = depth + nm.loc[1] = depth -# Get node depths -depths = dijkstra_shortest_paths(graph, 1).dists -nodes_per_depth = Dict(Int(depth) => Int[] for depth in unique(depths)) + push!(nodes_per_depth[Int(depth)], i) + end -for (i, depth) in enumerate(depths) - nm = graph[i] - nm.depth[] = depth - nm.loc[1] = depth + # Sort nodes by file for each depth + for nodes in values(nodes_per_depth) + sort!(nodes, by = i -> graph[i].file) + end - push!(nodes_per_depth[Int(depth)], i) + return nodes_per_depth end -# Squash per depth nodes with the same name into one -function squash!(graph, nodes_per_depth) - for depth in 1:maxdepth +function squash!(graph, nodes_per_depth, max_depth) + for depth in 1:max_depth names = Dict{String, Vector{Int}}() nodes_at_depth = nodes_per_depth[depth] for i in nodes_at_depth @@ -184,73 +164,120 @@ function squash!(graph, nodes_per_depth) end end -squash_per_depth && squash!(graph, nodes_per_depth) +function set_coordinates!(graph, nodes_per_depth, max_depth, plot_non_Ribasim) + for depth in 0:max_depth + nodes = nodes_per_depth[depth] + n_nodes = if plot_non_Ribasim + length(nodes) + else + count(i -> graph[i].mod == :Ribasim, nodes) + end + ys = n_nodes == 1 ? [0] : range(-1, 1, length = n_nodes) + idx = 1 -# Sort nodes by file for each depth -for nodes in values(nodes_per_depth) - sort!(nodes, by = i -> graph[i].file) + for i in nodes + nm = graph[i] + if (nm.mod == :Ribasim || plot_non_Ribasim) + graph[i].loc .= (depth, ys[idx]) + idx += 1 + end + end + end end -# Set node coordinates -for depth in 0:maxdepth - nodes = nodes_per_depth[depth] - n_nodes = if plot_non_Ribasim - length(nodes) - else - count(i -> graph[i].mod == :Ribasim, nodes) +function plot_edges!(ax, graph, max_depth) + for edge in edges(graph) + nm_src = graph[label_for(graph, edge.src)] + nm_dst = graph[label_for(graph, edge.dst)] + + (nm_dst.depth[] > max_depth) && continue + + x = [nm_src.loc[1], nm_dst.loc[1]] + y = [nm_src.loc[2], nm_dst.loc[2]] + lines!(ax, x, y; color = :black) end - ys = n_nodes == 1 ? [0] : range(-1, 1, length = n_nodes) - idx = 1 +end - for i in nodes - nm = graph[i] - if (nm.mod == :Ribasim || plot_non_Ribasim) - graph[i].loc .= (depth, ys[idx]) - idx += 1 - end +function plot_labels!(ax, graph, max_depth, color_dict) + for node in labels(graph) + nm = graph[node] + (nm.depth[] > max_depth) && continue + text!(ax, nm.loc..., + text = "$nm", + color = get(color_dict, nm.file, :black), + font = :bold, + strokecolor = :black, + strokewidth = 1, + label = String(nm.file) + ) end end -# Plot edges -for edge in edges(graph) - nm_src = graph[label_for(graph, edge.src)] - nm_dst = graph[label_for(graph, edge.dst)] +function plot_graph( + graph_orig::MetaGraph; + max_depth::Int = 5, + plot_non_Ribasim::Bool = false, + squash_per_depth::Bool = true, +) - (nm_dst.depth[] > maxdepth) && continue + f = Figure() + ax = Axis(f[1, 1]) - x = [nm_src.loc[1], nm_dst.loc[1]] - y = [nm_src.loc[2], nm_dst.loc[2]] - lines!(ax, x, y; color = :black) -end + graph = copy(graph_orig) + nodes_per_depth = get_node_depths(graph) + + # Squash per depth nodes with the same name into one + squash_per_depth && squash!(graph, nodes_per_depth, max_depth) + + set_coordinates!(graph, nodes_per_depth, max_depth, plot_non_Ribasim) + plot_edges!(ax, graph, max_depth) -# Plot node labels -files = sort(unique(graph[i].file for i in labels(graph) if graph[i].mod == :Ribasim), rev = true) -colors = distinguishable_colors(length(files)+1)[2:end] -color_dict = OrderedDict(zip(files, colors)) - -get_color(file) = get(color_dict, file, :black) - -for node in labels(graph) - nm = graph[node] - (nm.depth[] > maxdepth) && continue - color = nm.mod == :Ribasim ? :green : :black - text!(ax, nm.loc..., - text = "$nm", - color = get_color(nm.file), - font = :bold, - strokecolor = :black, - strokewidth = 1, - label = String(nm.file) + files = sort(unique(graph[i].file for i in labels(graph) if graph[i].mod == :Ribasim), rev = true) + colors = distinguishable_colors(length(files)+1)[2:end] + color_dict = OrderedDict(zip(files, colors)) + + plot_labels!(ax, graph, max_depth, color_dict) + + Legend( + f[1,2], + [MarkerElement(color = c, marker = :rect) for c in values(color_dict)], + String.(files) ) + + f end -Legend( - f[1,2], - [MarkerElement(color = c, marker = :rect) for c in values(color_dict)], - String.(files) -) +############################################# + +using Ribasim +``` -f +## Parameter initialization + +```{julia} +# | code-fold: true +using SQLite +toml_path = normpath(@__DIR__, "../../generated_testmodels/basic_transient/ribasim.toml") +config = Ribasim.Config(toml_path) +db_path = Ribasim.database_path(config) +db = SQLite.DB(db_path) + +graph, verts = tracecall((Ribasim,), Ribasim.Parameters, (db, config)) +close(db) + +plot_graph(graph) +``` + +## Water balance + +```{julia} +# | code-fold: true +using OrdinaryDiffEqCore: get_du +model = Ribasim.Model(toml_path) +du = get_du(model.integrator) +(; u, p, t) = model.integrator +graph, verts = tracecall((Ribasim,), Ribasim.water_balance!, (du, u, p, t)) +plot_graph(graph) ``` ```{julia}