Callstack visualization (#1922)

Fixes #1472.

Project.toml

@@ -7,7 +7,9 @@ Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
 Arrow = "69666777-d1a9-59fb-9406-91d4454c9d45"
 Artifacts = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"
 BasicModelInterface = "59605e27-edc0-445a-b93d-c09a3a50b330"
+CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 CodecZstd = "6b39b394-51ab-5f42-8807-6242bab2b4c2"
+Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 Configurations = "5218b696-f38b-4ac9-8b61-a12ec717816d"
 DBInterface = "a10d1c49-ce27-4219-8d33-6db1a4562965"
@@ -30,6 +32,7 @@ InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 IterTools = "c8e1da08-722c-5040-9ed9-7db0dc04731e"
 JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
 JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
+JuliaInterpreter = "aa1ae85d-cabe-5617-a682-6adf51b2e16a"
 Legolas = "741b9549-f6ed-4911-9fbf-4a1c0c97f0cd"
 LibGit2 = "76f85450-5226-5b5a-8eaa-529ad045b433"
 LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255"

core/src/model.jl

@@ -54,7 +54,7 @@ function Model(config::Config)::Model
         error("Invalid edge types found.")
     end
 
-    local parameters, state, n, tstops
+    local parameters, tstops
     try
         parameters = Parameters(db, config)
 

core/src/write.jl

@@ -136,7 +136,6 @@ function basin_table(
     data = get_storages_and_levels(model)
     storage = vec(data.storage[:, begin:(end - 1)])
     level = vec(data.level[:, begin:(end - 1)])
-    Δstorage = vec(diff(data.storage; dims = 2))
 
     nbasin = length(data.node_id)
     ntsteps = length(data.time) - 1

docs/.gitignore

@@ -5,4 +5,5 @@
 /tutorial/crystal-basin/
 guide/data/
 *.html
+*.quarto_ipynb
 objects.json

docs/_quarto.yml

@@ -98,6 +98,7 @@ website:
         - dev/bmi.qmd
         - dev/ci.qmd
         - dev/release.qmd
+        - dev/callstacks.qmd
         - section: QGIS
           contents:
             - dev/qgis.qmd

docs/concept/equations.qmd

@@ -105,7 +105,7 @@ Hence the reference used for computing the relative error is the average of the
 
 The default tolerances are $0.001 \text{ m}^3$ for the balance error and $0.01$ for the relative error, which should not be exceeded for realistic models.
 
-In extreme cases where the storage rate is many orders of magnitude smaller than the storage itself, these computations can have floating point truncation errors which can lead to large relative errors.  This is however only when the storage is roughly $ \ge 10^{15}$ times bigger than the storage rate.
+In extreme cases where the storage rate is many orders of magnitude smaller than the storage itself, these computations can have floating point truncation errors which can lead to large relative errors.  This is however only when the storage is roughly $\geq 10^{15}$ times bigger than the storage rate.
 
 ### Example calculation
 

docs/dev/callstacks.qmd

@@ -0,0 +1,115 @@
+# Call stacks
+
+```{julia}
+# | code-fold: true
+using CairoMakie
+using Colors
+using Graphs, MetaGraphsNext
+using JuliaInterpreter, OrderedCollections
+
+include("scripts/trace_call.jl")
+include("scripts/plot_trace.jl")
+
+using Ribasim
+using Random
+
+Random.seed!(1);
+```
+
+The plots below show the call stack within the Julia core for several important entrypoints. The function names are colored by the script in which they are defined, and the lines between the function names have random colors to be able to differentiate between them. Solid lines refer to calls to functions defined in the same script, dashed ones to functions defined in a different script. The plots are of high resolution so zooming in to particular parts is encouraged.
+
+Note that these graphs are obtained by dynamic analysis, i.e. by running parts of the code with specific inputs. This means that there can be unshown paths through the code that are not reached for these particular inputs.
+
+## Parameter initialization
+
+Parameter initialization is the process of reading the parameter values from the input files and storing them in data structures for easy access in the core. Most notable here is the convergence of many paths to `load_structvector` and `parse_static_and_time`, as these are needed for parameter initialization for most node types.
+
+```{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;
+    size = (2000, 1200),
+    squash_methods = [
+        :n_neighbor_bounds_flow,
+        :n_neighbor_bounds_control,
+        :sort_by_function,
+        :neighbortypes
+    ],
+    xlims = (-0.4, 5.6)
+)
+```
+
+## `water_balance!`
+
+`water_balance!` is the right hand side function of the system of ODEs that is solved by the Ribasim core (for more details see [here](../concept/equations.qmd#formal-model-description)). The various `formulate_flow!` methods are for flow rates as determined by different node types.
+
+```{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, size = (1700, 1000), xlims = (-0.4, 4.5))
+```
+
+## Allocation initialization
+
+In this part of the code the data structures for allocation are set up. Most endpoints in `allocation_init.jl` set up data structures as defined in [JuMP.jl](https://jump.dev/JuMP.jl/stable/).
+
+```{julia}
+# | code-fold: true
+toml_path = normpath(@__DIR__, "../../generated_testmodels/main_network_with_subnetworks/ribasim.toml")
+config = Ribasim.Config(toml_path; allocation_use_allocation=false)
+db_path = Ribasim.database_path(config)
+db = SQLite.DB(db_path)
+p = Ribasim.Parameters(db, config)
+graph, verts = tracecall((Ribasim,), Ribasim.initialize_allocation!, (p, config))
+plot_graph(graph, size = (1800, 1000), xlims = (-0.5, 5.5))
+```
+
+## Allocation run
+
+Running the allocation algorithm consists of running the optimization itself (which is handled in `JuMP.jl`), and all Ribasim functions around it are used for communicating data between the optimization problem and the physical layer, as well as gathering output data. Fore more information on the allocation algorithm see [here](../concept/allocation.qmd).
+
+
+```{julia}
+# | code-fold: true
+model = Ribasim.Model(toml_path)
+graph, verts = tracecall((Ribasim,), Ribasim.update_allocation!, (model.integrator,))
+plot_graph(graph, size = (2000, 1000), xlims = (-0.4, 5.5))
+```
+
+## Discrete control
+
+Discrete control works by a [`FunctionCallingCallback`](https://docs.sciml.ai/DiffEqCallbacks/stable/output_saving/#DiffEqCallbacks.FunctionCallingCallback), changing parameters when a change in control state is detected (see also [here](../reference/node/discrete-control.qmd)).
+
+```{julia}
+# | code-fold: true
+toml_path = normpath(@__DIR__, "../../generated_testmodels/pump_discrete_control/ribasim.toml")
+model = Ribasim.Model(toml_path)
+(; u, t) = model.integrator
+model.integrator.p.basin.storage0 .= [0.1, 100.0]
+graph, verts = tracecall((Ribasim,), Ribasim.apply_discrete_control!, (u, t, model.integrator))
+plot_graph(graph; size = (1300, 500), prune_from = [:water_balance!], xlims = (-0.5, 3.5))
+```
+
+## Writing output
+
+Writing output (currently) happens only after the full simulation is finished. For more information on the different output tables see [here](../reference/usage.qmd#results).
+```{julia}
+# | code-fold: true
+toml_path = normpath(@__DIR__, "../../generated_testmodels/basic_transient/ribasim.toml")
+model = Ribasim.Model(toml_path)
+graph, verts = tracecall((Ribasim,), Ribasim.write_results, (model,))
+plot_graph(graph, size = (1600, 1000), xlims = (-0.5, 4.5))
+```

docs/dev/scripts/plot_trace.jl

@@ -0,0 +1,215 @@
+function cut_generated_calls!(graph)
+    for i in collect(labels(graph))
+        nm = graph[i]
+        (; name) = nm
+        if startswith(String(name), "#")
+            for i_in in inneighbor_labels(graph, i)
+                for i_out in outneighbor_labels(graph, i)
+                    graph[i_in, i_out] = nothing
+                end
+            end
+            delete!(graph, i)
+        end
+    end
+end
+
+function get_node_depths(graph)
+    depths = dijkstra_shortest_paths(graph, 1).dists
+    nodes_per_depth = Dict(Int(depth) => Int[] for depth in unique(depths))
+
+    for (i, depth) in zip(labels(graph), depths)
+        nm = graph[i]
+        nm.depth[] = depth
+        nm.loc[1] = depth
+
+        push!(nodes_per_depth[Int(depth)], i)
+    end
+
+    # Sort nodes by file for each depth
+    for nodes in values(nodes_per_depth)
+        sort!(nodes; by = i -> graph[i].file, rev = true)
+    end
+
+    return nodes_per_depth
+end
+
+function prune_branch!(
+    graph,
+    start::Int;
+    branch_base::Bool = true,
+    to_delete::Vector{Int} = Int[],
+)
+    branch_base && empty!(to_delete)
+    for i in outneighbor_labels(graph, start)
+        prune_branch!(graph, i; branch_base = false, to_delete)
+        push!(to_delete, i)
+    end
+    branch_base && delete!.(Ref(graph), to_delete)
+end
+
+function squash!(graph, nodes_per_depth, max_depth, squash_methods)
+    for depth in 1:max_depth
+        names = Dict{String, Vector{Int}}()
+        nodes_at_depth = nodes_per_depth[depth]
+        for i in nodes_at_depth
+            nm = graph[i]
+            name = if nm.name in squash_methods
+                "$(nm.mod).$(nm.name)"
+            else
+                "$nm"
+            end
+            if name in keys(names)
+                push!(names[name], i)
+            else
+                names[name] = [i]
+            end
+        end
+        for nodes in values(names)
+            (length(nodes) == 1) && continue
+            survivor = first(nodes)
+
+            for i in nodes[2:end]
+                for i_in in inneighbor_labels(graph, i)
+                    graph[i_in, survivor] = nothing
+                    delete!(graph, i_in, i)
+                end
+
+                for i_out in outneighbor_labels(graph, i)
+                    graph[survivor, i_out] = nothing
+                    delete!(graph, i, i_out)
+                end
+
+                delete!(graph, i)
+                deleteat!(nodes_at_depth, findfirst(==(i), nodes_at_depth))
+            end
+        end
+    end
+end
+
+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.5] : range(0, 1; length = n_nodes)
+        idx = 1
+
+        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
+
+function plot_edges!(ax, graph, max_depth, nodes_per_depth; n_points = 25)
+    for depth in 0:(max_depth - 1)
+        nodes_at_depth = nodes_per_depth[depth]
+        n_nodes = length(nodes_at_depth)
+        for (idx, i) in enumerate(nodes_at_depth)
+            nm_src = graph[i]
+            for i_out in outneighbor_labels(graph, i)
+                nm_dst = graph[i_out]
+
+                A = (nm_src.loc[2] - nm_dst.loc[2]) / 2
+                B = π / (nm_dst.loc[1] - nm_src.loc[1])
+                C = (nm_src.loc[2] + nm_dst.loc[2]) / 2
+
+                x = range(nm_src.loc[1], nm_dst.loc[1]; length = n_points)
+                y = @. A * cos(B * (x - nm_src.loc[1])) + C
+
+                color = RGB((0.8 * rand(3))...)
+                linestyle = (nm_src.file == nm_dst.file) ? :solid : :dash
+                lines!(ax, x, y; color, linestyle)
+            end
+        end
+    end
+end
+
+function plot_labels!(ax, graph, max_depth, color_dict)
+    for node in labels(graph)
+        nm = graph[node]
+        x, y = nm.loc
+        (nm.depth[] > max_depth) && continue
+        text!(
+            ax,
+            x,
+            y;
+            text = "$nm",
+            color = get(color_dict, nm.file, :black),
+            font = :bold,
+            strokecolor = :black,
+            strokewidth = 1.0,
+            label = String(nm.file),
+            align = (:center, :bottom),
+        )
+        scatter!(ax, [x], [y]; color = :black)
+    end
+end
+
+function plot_graph(
+    graph_orig::MetaGraph;
+    size = (1000, 1000),
+    max_depth::Int = 5,
+    plot_non_Ribasim::Bool = false,
+    squash_per_depth::Bool = true,
+    squash_methods::Vector{Symbol} = Symbol[],
+    prune_from::Vector{Symbol} = Symbol[],
+    xlims = nothing,
+)
+    graph = copy(graph_orig)
+
+    # Prune branches
+    for i in collect(labels(graph))
+        if haskey(graph, i)
+            nm = graph[i]
+            if nm.name in prune_from
+                prune_branch!(graph, i)
+            end
+        end
+    end
+
+    # Cut out calls whose name starts with '#'
+    cut_generated_calls!(graph)
+
+    nodes_per_depth = get_node_depths(graph)
+    max_depth = min(max_depth, maximum(keys(nodes_per_depth)))
+
+    # Squash per depth nodes with the same name into one
+    squash_per_depth && squash!(graph, nodes_per_depth, max_depth, squash_methods)
+
+    set_coordinates!(graph, nodes_per_depth, max_depth, plot_non_Ribasim)
+
+    files = sort(unique(graph[i].file for i in labels(graph) if graph[i].mod == :Ribasim))
+    colors = distinguishable_colors(length(files) + 1)[end:-1:2]
+    color_dict = OrderedDict(zip(files, colors))
+
+    delete!(theme(nothing), :resolution) # Needed because of a refactor in Makie going from resolution to size
+    f = Figure(; size = size)
+    ax = Axis(f[1, 1]; xlabel = "depth", xticks = 0:max_depth)
+    plot_edges!(ax, graph, max_depth, nodes_per_depth)
+    plot_labels!(ax, graph, max_depth, color_dict)
+    hideydecorations!(ax)
+    !isnothing(xlims) && xlims!(ax, xlims...)
+
+    # Build legend
+    elements = LegendElement[
+        MarkerElement(; color = c, marker = :rect) for c in values(color_dict)
+    ]
+    descriptions = String.(files)
+
+    push!(elements, LineElement(; color = :black, linestyle = :dash))
+    push!(descriptions, "between scripts")
+
+    push!(elements, LineElement(; color = :black, linestyle = :solid))
+    push!(descriptions, "within a script")
+
+    Legend(f[1, 2], elements, descriptions)
+
+    f
+end