diff --git a/docs/dev/callstacks.qmd b/docs/dev/callstacks.qmd
index ccad1c637..5bd4fe372 100644
--- a/docs/dev/callstacks.qmd
+++ b/docs/dev/callstacks.qmd
@@ -1,3 +1,7 @@
+---
+lightbox: auto
+---
+
 # Call stacks
 
 ```{julia}
@@ -37,14 +41,12 @@ close(db)
 
 plot_graph(
     graph;
-    size = (2000, 1200),
     squash_methods = [
         :n_neighbor_bounds_flow,
         :n_neighbor_bounds_control,
         :sort_by_function,
-        :neighbortypes
+        :neighbortypes,
     ],
-    xlims = (-0.4, 5.6)
 )
 ```
 
@@ -59,7 +61,7 @@ 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))
+plot_graph(graph; max_depth = 4)
 ```
 
 ## Allocation initialization
@@ -68,13 +70,16 @@ In this part of the code the data structures for allocation are set up. Most end
 
 ```{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)
+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))
+plot_graph(graph)
 ```
 
 ## Allocation run
@@ -86,7 +91,7 @@ Running the allocation algorithm consists of running the optimization itself (wh
 # | 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))
+plot_graph(graph)
 ```
 
 ## Discrete control
@@ -95,12 +100,14 @@ Discrete control works by a [`FunctionCallingCallback`](https://docs.sciml.ai/Di
 
 ```{julia}
 # | code-fold: true
-toml_path = normpath(@__DIR__, "../../generated_testmodels/pump_discrete_control/ribasim.toml")
+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))
+graph, verts =
+    tracecall((Ribasim,), Ribasim.apply_discrete_control!, (u, t, model.integrator))
+plot_graph(graph; prune_from = [:water_balance!], max_depth = 3)
 ```
 
 ## Writing output
@@ -111,5 +118,5 @@ Writing output (currently) happens only after the full simulation is finished. F
 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))
+plot_graph(graph; max_depth = 4)
 ```
diff --git a/docs/dev/scripts/plot_trace.jl b/docs/dev/scripts/plot_trace.jl
index 085a04a71..bd966b2bb 100644
--- a/docs/dev/scripts/plot_trace.jl
+++ b/docs/dev/scripts/plot_trace.jl
@@ -123,7 +123,7 @@ function plot_edges!(ax, graph, max_depth, nodes_per_depth; n_points = 25)
                 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))...)
+                color = RGBA((0.8 * rand(3))..., 0.5)
                 linestyle = (nm_src.file == nm_dst.file) ? :solid : :dash
                 lines!(ax, x, y; color, linestyle)
             end
@@ -141,10 +141,10 @@ function plot_labels!(ax, graph, max_depth, color_dict)
             x,
             y;
             text = "$nm",
-            color = get(color_dict, nm.file, :black),
+            color = :black,
             font = :bold,
-            strokecolor = :black,
-            strokewidth = 1.0,
+            strokecolor = get(color_dict, nm.file, :black),
+            strokewidth = 0.5,
             label = String(nm.file),
             align = (:center, :bottom),
         )
@@ -154,7 +154,7 @@ end
 
 function plot_graph(
     graph_orig::MetaGraph;
-    size = (1000, 1000),
+    size = (2000, 1000),
     max_depth::Int = 5,
     plot_non_Ribasim::Bool = false,
     squash_per_depth::Bool = true,
@@ -189,19 +189,22 @@ function plot_graph(
     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
+    theme = theme_minimal()
+    set_theme!(theme)
+    delete!(theme, :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)
+    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...)
+    hidespines!(ax)
+    isnothing(xlims) ? xlims!(ax, -0.25, max_depth + 0.5) : xlims!(ax, xlims...)
 
     # Build legend
     elements = LegendElement[
         MarkerElement(; color = c, marker = :rect) for c in values(color_dict)
     ]
-    descriptions = String.(files)
+    descriptions = basename.(String.(files))
 
     push!(elements, LineElement(; color = :black, linestyle = :dash))
     push!(descriptions, "between scripts")
@@ -209,7 +212,16 @@ function plot_graph(
     push!(elements, LineElement(; color = :black, linestyle = :solid))
     push!(descriptions, "within a script")
 
-    Legend(f[1, 2], elements, descriptions)
-
+    axislegend(
+        ax,
+        elements,
+        descriptions;
+        position = :lt,
+        framevisible = true,
+        margin = (20, 20, 20, 20),
+        padding = 10,
+        framecolor = :lightgrey,
+    )
+    resize_to_layout!(f)
     f
 end