diff --git a/Manifest.toml b/Manifest.toml
index 4ce131ea4..5124ce6a6 100644
--- a/Manifest.toml
+++ b/Manifest.toml
@@ -2,7 +2,7 @@
 
 julia_version = "1.11.1"
 manifest_format = "2.0"
-project_hash = "921c2f0c042509feb4a1950ec24084c1ad561dd0"
+project_hash = "ecb70ca5008f6b49d97f55488305bbc1ef31ed96"
 
 [[deps.ADTypes]]
 git-tree-sha1 = "eea5d80188827b35333801ef97a40c2ed653b081"
diff --git a/Project.toml b/Project.toml
index 1c5d51394..ec73afc80 100644
--- a/Project.toml
+++ b/Project.toml
@@ -9,6 +9,7 @@ Artifacts = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"
 BasicModelInterface = "59605e27-edc0-445a-b93d-c09a3a50b330"
 CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 CodecZstd = "6b39b394-51ab-5f42-8807-6242bab2b4c2"
+ColorSchemes = "35d6a980-a343-548e-a6ea-1d62b119f2f4"
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 Configurations = "5218b696-f38b-4ac9-8b61-a12ec717816d"
diff --git a/docs/dev/callstacks.qmd b/docs/dev/callstacks.qmd
index cdcf0bf1c..d7210984a 100644
--- a/docs/dev/callstacks.qmd
+++ b/docs/dev/callstacks.qmd
@@ -192,7 +192,7 @@ function set_coordinates!(graph, nodes_per_depth, max_depth, plot_non_Ribasim)
         else
             count(i -> graph[i].mod == :Ribasim, nodes)
         end
-        ys = n_nodes == 1 ? [0] : range(-1, 1, length = n_nodes)
+        ys = n_nodes == 1 ? [0.5] : range(0, 1, length = n_nodes)
         idx = 1
 
         for i in nodes
@@ -205,21 +205,28 @@ function set_coordinates!(graph, nodes_per_depth, max_depth, plot_non_Ribasim)
     end
 end
 
-function plot_edges!(ax, graph, max_depth, n_points = 25)
-    for edge in edges(graph)
-        nm_src = graph[label_for(graph, edge.src)]
-        nm_dst = graph[label_for(graph, edge.dst)]
+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]
 
-        (nm_dst.depth[] > max_depth) && continue
+                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
 
-        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
 
-        x = range(nm_src.loc[1], nm_dst.loc[1], length = n_points)
-        y = @. A * cos(B * (x - nm_src.loc[1])) + C
 
-        lines!(ax, x, y; color = :black)
+                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
 
@@ -250,7 +257,7 @@ function plot_graph(
     squash_methods::Vector{Symbol} = Symbol[],
 )
     delete!(theme(nothing), :resolution) # Needed because of a refactor in Makie going from resolution to size
-    f = Figure(; size = size)
+    f = Figure(; size = size, xlabel = "depth")
     ax = Axis(f[1, 1])
 
     graph = copy(graph_orig)
@@ -259,12 +266,13 @@ function plot_graph(
     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)
-    plot_edges!(ax, graph, max_depth)
+    plot_edges!(ax, graph, max_depth, nodes_per_depth)
 
     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]
@@ -272,10 +280,20 @@ function plot_graph(
 
     plot_labels!(ax, graph, max_depth, color_dict)
 
+    # 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],
-        [MarkerElement(color = c, marker = :rect) for c in values(color_dict)],
-        String.(files)
+        elements,
+        descriptions
     )
 
     f
@@ -284,6 +302,9 @@ end
 #############################################
 
 using Ribasim
+using Random
+
+Random.seed!(1)
 ```
 
 ## Parameter initialization
@@ -301,7 +322,7 @@ close(db)
 
 plot_graph(
     graph;
-    size = (1500, 1200),
+    size = (1800, 1200),
     squash_methods = [
         :n_neighbor_bounds_flow,
         :n_neighbor_bounds_control,
@@ -320,14 +341,21 @@ 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)
+plot_graph(graph, size = (1600, 1000))
 ```
 
+## Allocation initialization
 ```{julia}
-# TODO:
-# - Fix y coordinate of nodes that have no outneighbors
-# - Remove nodes (and everything downstream) from non-Ribasim calls
-# - Color boxes by file name and add legend
-# - Add boxes around names
-# - Let connection lines emerge from box boundaries
+# | code-fold: true
+toml_path = normpath(@__DIR__, "../../generated_testmodels/main_network_with_subnetworks/ribasim.toml")
+model = Ribasim.Model(toml_path)
+# graph, verts = tracecall((Ribasim,), Ribasim.initialize_allocation!, (model.integrator.p, model.config))
+# plot_graph(graph)
+```
+
+## Allocation run
+```{julia}
+# | code-fold: true
+graph, verts = tracecall((Ribasim,), Ribasim.update_allocation!, (model.integrator,))
+plot_graph(graph, size = (1600, 1000))
 ```