diff --git a/Manifest.toml b/Manifest.toml
index c49980a83..13ade21fa 100644
--- a/Manifest.toml
+++ b/Manifest.toml
@@ -2,7 +2,7 @@
 
 julia_version = "1.10.5"
 manifest_format = "2.0"
-project_hash = "0257f2772e4bfed0b6316b6871c4495978c173b4"
+project_hash = "a2d1a982c3293971ae40e0a7bdfb40a85bd30ac1"
 
 [[deps.ADTypes]]
 git-tree-sha1 = "eea5d80188827b35333801ef97a40c2ed653b081"
@@ -1346,7 +1346,7 @@ uuid = "295af30f-e4ad-537b-8983-00126c2a3abe"
 version = "3.6.0"
 
 [[deps.Ribasim]]
-deps = ["Accessors", "Arrow", "BasicModelInterface", "CodecZstd", "ComponentArrays", "Configurations", "DBInterface", "DataInterpolations", "DataStructures", "Dates", "DiffEqBase", "DiffEqCallbacks", "EnumX", "FiniteDiff", "Graphs", "HiGHS", "IterTools", "JuMP", "Legolas", "LineSearches", "LinearSolve", "Logging", "LoggingExtras", "MetaGraphsNext", "OrdinaryDiffEqBDF", "OrdinaryDiffEqCore", "OrdinaryDiffEqLowOrderRK", "OrdinaryDiffEqNonlinearSolve", "OrdinaryDiffEqRosenbrock", "OrdinaryDiffEqSDIRK", "OrdinaryDiffEqTsit5", "PreallocationTools", "SQLite", "SciMLBase", "SparseArrays", "SparseConnectivityTracer", "StructArrays", "Tables", "TerminalLoggers", "TranscodingStreams"]
+deps = ["Accessors", "Arrow", "BasicModelInterface", "CodecZstd", "ComponentArrays", "Configurations", "DBInterface", "DataInterpolations", "DataStructures", "Dates", "DiffEqBase", "DiffEqCallbacks", "EnumX", "FiniteDiff", "Graphs", "HiGHS", "IterTools", "JuMP", "Legolas", "LineSearches", "LinearAlgebra", "LinearSolve", "Logging", "LoggingExtras", "MetaGraphsNext", "OrdinaryDiffEqBDF", "OrdinaryDiffEqCore", "OrdinaryDiffEqLowOrderRK", "OrdinaryDiffEqNonlinearSolve", "OrdinaryDiffEqRosenbrock", "OrdinaryDiffEqSDIRK", "OrdinaryDiffEqTsit5", "PreallocationTools", "SQLite", "SciMLBase", "SparseArrays", "SparseConnectivityTracer", "StructArrays", "Tables", "TerminalLoggers", "TranscodingStreams"]
 path = "core"
 uuid = "aac5e3d9-0b8f-4d4f-8241-b1a7a9632635"
 version = "2024.11.0"
diff --git a/core/Project.toml b/core/Project.toml
index fd61f6262..196370cc5 100644
--- a/core/Project.toml
+++ b/core/Project.toml
@@ -25,6 +25,7 @@ IterTools = "c8e1da08-722c-5040-9ed9-7db0dc04731e"
 JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
 Legolas = "741b9549-f6ed-4911-9fbf-4a1c0c97f0cd"
 LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 LoggingExtras = "e6f89c97-d47a-5376-807f-9c37f3926c36"
@@ -77,6 +78,7 @@ IOCapture = "0.2"
 IterTools = "1.4"
 JuMP = "1.15"
 Legolas = "0.5"
+LinearAlgebra = "1"
 LineSearches = "7"
 LinearSolve = "2.24"
 Logging = "<0.0.1, 1"
diff --git a/core/src/Ribasim.jl b/core/src/Ribasim.jl
index 0c0ab18f5..c95296756 100644
--- a/core/src/Ribasim.jl
+++ b/core/src/Ribasim.jl
@@ -46,6 +46,12 @@ using SciMLBase:
 # through operator overloading
 using SparseConnectivityTracer: TracerSparsityDetector, jacobian_sparsity, GradientTracer
 
+# For efficient sparse computations
+using SparseArrays: SparseMatrixCSC, spzeros
+
+# Linear algebra
+using LinearAlgebra: mul!
+
 # PreallocationTools is used because the RHS function (water_balance!) gets called with different input types
 # for u, du:
 # - Float64 for normal calls
@@ -92,7 +98,7 @@ using TerminalLoggers: TerminalLogger
 # Convenience wrapper around arrays, divides vectors in
 # separate sections which can be indexed individually.
 # Used for e.g. Basin forcing and the state vector.
-using ComponentArrays: ComponentVector, Axis
+using ComponentArrays: ComponentVector, ComponentArray, Axis, getaxes
 
 # Date and time handling; externally we use the proleptic Gregorian calendar,
 # internally we use a Float64; seconds since the start of the simulation.
diff --git a/core/src/callback.jl b/core/src/callback.jl
index 0cdd0d23b..712cb868a 100644
--- a/core/src/callback.jl
+++ b/core/src/callback.jl
@@ -256,20 +256,23 @@ function save_flow(u, t, integrator)
         drainage,
         t,
     )
-    check_water_balance_error(integrator, saved_flow, Δt)
+    check_water_balance_error!(saved_flow, integrator, Δt)
     return saved_flow
 end
 
-function check_water_balance_error(
-    integrator::DEIntegrator,
+function check_water_balance_error!(
     saved_flow::SavedFlow,
+    integrator::DEIntegrator,
     Δt::Float64,
 )::Nothing
     (; u, p, t) = integrator
     (; basin, water_balance_abstol, water_balance_reltol) = p
     errors = false
     current_storage = basin.current_storage[parent(u)]
-    formulate_storages!(current_storage, u, u, p, t)
+
+    # The initial storage is irrelevant for the storage rate and can only cause
+    # floating point truncation errors
+    formulate_storages!(current_storage, u, u, p, t; add_initial_storage = false)
 
     for (
         inflow_rate,
@@ -289,7 +292,7 @@ function check_water_balance_error(
         saved_flow.flow.evaporation,
         saved_flow.flow.infiltration,
         current_storage,
-        basin.storage_prev_saveat,
+        basin.Δstorage_prev_saveat,
         basin.node_id,
     )
         storage_rate = (s_now - s_prev) / Δt
@@ -304,13 +307,17 @@ function check_water_balance_error(
             errors = true
             @error "Too large water balance error" id balance_error relative_error
         end
+
+        saved_flow.storage_rate[id.idx] = storage_rate
+        saved_flow.balance_error[id.idx] = balance_error
+        saved_flow.relative_error[id.idx] = relative_error
     end
     if errors
         t = datetime_since(t, p.starttime)
         error("Too large water balance error(s) detected at t = $t")
     end
 
-    @. basin.storage_prev_saveat = current_storage
+    @. basin.Δstorage_prev_saveat = current_storage
     return nothing
 end
 
diff --git a/core/src/model.jl b/core/src/model.jl
index 148f0bc5d..d4aa9af1e 100644
--- a/core/src/model.jl
+++ b/core/src/model.jl
@@ -104,6 +104,7 @@ function Model(config::Config)::Model
     du0 = zero(u0)
 
     parameters = set_state_flow_edges(parameters, u0)
+    parameters = build_flow_to_storage(parameters, u0)
     parameters = @set parameters.u_prev_saveat = zero(u0)
 
     # The Solver algorithm
diff --git a/core/src/parameter.jl b/core/src/parameter.jl
index 576c12769..ac5ded66e 100644
--- a/core/src/parameter.jl
+++ b/core/src/parameter.jl
@@ -268,6 +268,8 @@ In-memory storage of saved mean flows for writing to results.
 - `flow_boundary`: The exact integrated mean flows of flow boundaries
 - `precipitation`: The exact integrated mean precipitation
 - `drainage`: The exact integrated mean drainage
+- `balance_error`: The (absolute) water balance error
+- `relative_error`: The relative water balance error
 - `t`: Endtime of the interval over which is averaged
 """
 @kwdef struct SavedFlow{V}
@@ -277,6 +279,9 @@ In-memory storage of saved mean flows for writing to results.
     flow_boundary::Vector{Float64}
     precipitation::Vector{Float64}
     drainage::Vector{Float64}
+    storage_rate::Vector{Float64} = zero(precipitation)
+    balance_error::Vector{Float64} = zero(precipitation)
+    relative_error::Vector{Float64} = zero(precipitation)
     t::Float64
 end
 
@@ -313,7 +318,8 @@ end
     vertical_flux::V = zeros(length(node_id))
     # Initial_storage
     storage0::Vector{Float64} = zeros(length(node_id))
-    storage_prev_saveat::Vector{Float64} = zeros(length(node_id))
+    # Storage at previous saveat without storage0
+    Δstorage_prev_saveat::Vector{Float64} = zeros(length(node_id))
     # Analytically integrated forcings
     cumulative_precipitation::Vector{Float64} = zeros(length(node_id))
     cumulative_drainage::Vector{Float64} = zeros(length(node_id))
@@ -840,6 +846,8 @@ const ModelGraph = MetaGraph{
     state_outflow_edge::C4 = ComponentVector()
     all_nodes_active::Base.RefValue{Bool} = Ref(false)
     tprev::Base.RefValue{Float64} = Ref(0.0)
+    # Sparse matrix for combining flows into storages
+    flow_to_storage::SparseMatrixCSC{Float64, Int64} = spzeros(1, 1)
     # Water balance tolerances
     water_balance_abstol::Float64
     water_balance_reltol::Float64
diff --git a/core/src/read.jl b/core/src/read.jl
index 7983a3e5e..bb259c3c3 100644
--- a/core/src/read.jl
+++ b/core/src/read.jl
@@ -649,7 +649,6 @@ function Basin(db::DB, config::Config, graph::MetaGraph)::Basin
     storage0 = get_storages_from_levels(basin, state.level)
     @assert length(storage0) == n "Basin / state length differs from number of Basins"
     basin.storage0 .= storage0
-    basin.storage_prev_saveat .= storage0
     return basin
 end
 
diff --git a/core/src/solve.jl b/core/src/solve.jl
index 2b0ca0b4e..9c2555908 100644
--- a/core/src/solve.jl
+++ b/core/src/solve.jl
@@ -118,86 +118,40 @@ function formulate_storages!(
     du::ComponentVector,
     u::ComponentVector,
     p::Parameters,
-    t::Number,
+    t::Number;
+    add_initial_storage::Bool = true,
 )::Nothing
-    (;
-        basin,
-        flow_boundary,
-        tabulated_rating_curve,
-        pump,
-        outlet,
-        linear_resistance,
-        manning_resistance,
-        user_demand,
-        tprev,
-    ) = p
-    # Current storage: initial conditdion +
+    (; basin, flow_boundary, tprev, flow_to_storage) = p
+    # Current storage: initial condition +
     # total inflows and outflows since the start
     # of the simulation
-    current_storage .= basin.storage0
-    formulate_storage!(current_storage, basin, du, u)
+    if add_initial_storage
+        current_storage .= basin.storage0
+    else
+        current_storage .= 0.0
+    end
+    mul!(current_storage, flow_to_storage, u, 1, 1)
+    formulate_storage!(current_storage, basin, du)
     formulate_storage!(current_storage, tprev[], t, flow_boundary)
-    formulate_storage!(current_storage, t, u.tabulated_rating_curve, tabulated_rating_curve)
-    formulate_storage!(current_storage, t, u.pump, pump)
-    formulate_storage!(current_storage, t, u.outlet, outlet)
-    formulate_storage!(current_storage, t, u.linear_resistance, linear_resistance)
-    formulate_storage!(current_storage, t, u.manning_resistance, manning_resistance)
-    formulate_storage!(
-        current_storage,
-        t,
-        u.user_demand_inflow,
-        user_demand;
-        edge_volume_out = u.user_demand_outflow,
-    )
     return nothing
 end
 
 """
-The storage contributions of the forcings that are part of the state.
+The storage contributions of the forcings that are not part of the state.
 """
 function formulate_storage!(
     current_storage::AbstractVector,
     basin::Basin,
     du::ComponentVector,
-    u::ComponentVector,
 )
     (; current_cumulative_precipitation, current_cumulative_drainage) = basin
 
-    current_storage .-= u.evaporation
-    current_storage .-= u.infiltration
-
     current_cumulative_precipitation = current_cumulative_precipitation[parent(du)]
     current_cumulative_drainage = current_cumulative_drainage[parent(du)]
     current_storage .+= current_cumulative_precipitation
     current_storage .+= current_cumulative_drainage
 end
 
-"""
-Formulate storage contributions of nodes.
-"""
-function formulate_storage!(
-    current_storage::AbstractVector,
-    t::Number,
-    edge_volume_in::AbstractVector,
-    node::AbstractParameterNode;
-    edge_volume_out = nothing,
-)
-    edge_volume_out = isnothing(edge_volume_out) ? edge_volume_in : edge_volume_out
-
-    for (volume_in, volume_out, inflow_edge, outflow_edge) in
-        zip(edge_volume_in, edge_volume_out, node.inflow_edge, node.outflow_edge)
-        inflow_id = inflow_edge.edge[1]
-        if inflow_id.type == NodeType.Basin
-            current_storage[inflow_id.idx] -= volume_in
-        end
-
-        outflow_id = outflow_edge.edge[2]
-        if outflow_id.type == NodeType.Basin
-            current_storage[outflow_id.idx] += volume_out
-        end
-    end
-end
-
 """
 Formulate storage contributions of flow boundaries.
 """
diff --git a/core/src/util.jl b/core/src/util.jl
index 38bed6f29..f8f17c8c0 100644
--- a/core/src/util.jl
+++ b/core/src/util.jl
@@ -641,7 +641,7 @@ function get_variable_ref(
     variable::String;
     listen::Bool = true,
 )::Tuple{PreallocationRef, Bool}
-    (; basin, graph) = p
+    (; basin) = p
     errors = false
 
     # Only built here because it is needed to obtain indices
@@ -840,7 +840,8 @@ but the derivative is bounded at x = 0.
 """
 function relaxed_root(x, threshold)
     if abs(x) < threshold
-        1 / 4 * (x / sqrt(threshold)) * (5 - (x / threshold)^2)
+        x_scaled = x / threshold
+        sqrt(threshold) * x_scaled^3 * (9 - 5x_scaled^2) / 4
     else
         sign(x) * sqrt(abs(x))
     end
@@ -953,6 +954,56 @@ function build_state_vector(p::Parameters)
     )
 end
 
+function build_flow_to_storage(p::Parameters, u::ComponentVector)::Parameters
+    n_basins = length(p.basin.node_id)
+    n_states = length(u)
+    flow_to_storage = ComponentArray(
+        spzeros(n_basins, n_states),
+        (Axis(; basins = 1:n_basins), only(getaxes(u))),
+    )
+
+    for node_name in (
+        :tabulated_rating_curve,
+        :pump,
+        :outlet,
+        :linear_resistance,
+        :manning_resistance,
+        :user_demand,
+    )
+        node = getfield(p, node_name)
+
+        if node_name == :user_demand
+            flow_to_storage_node_inflow = view(flow_to_storage, :, :user_demand_inflow)
+            flow_to_storage_node_outflow = view(flow_to_storage, :, :user_demand_outflow)
+        else
+            flow_to_storage_node_inflow = view(flow_to_storage, :, node_name)
+            flow_to_storage_node_outflow = flow_to_storage_node_inflow
+        end
+
+        for (inflow_edge, outflow_edge) in zip(node.inflow_edge, node.outflow_edge)
+            inflow_id, node_id = inflow_edge.edge
+            if inflow_id.type == NodeType.Basin
+                flow_to_storage_node_inflow[inflow_id.idx, node_id.idx] = -1.0
+            end
+
+            outflow_id = outflow_edge.edge[2]
+            if outflow_id.type == NodeType.Basin
+                flow_to_storage_node_outflow[outflow_id.idx, node_id.idx] = 1.0
+            end
+        end
+    end
+
+    flow_to_storage_evaporation = view(flow_to_storage, :, :evaporation)
+    flow_to_storage_infiltration = view(flow_to_storage, :, :infiltration)
+
+    for i in 1:n_basins
+        flow_to_storage_evaporation[i, i] = -1.0
+        flow_to_storage_infiltration[i, i] = -1.0
+    end
+
+    @set p.flow_to_storage = parent(flow_to_storage)
+end
+
 """
 Create vectors state_inflow_edge and state_outflow_edge which give for each state
 in the state vector in order the metadata of the edge that is associated with that state.
diff --git a/core/src/write.jl b/core/src/write.jl
index eec246653..1c0e1eeef 100644
--- a/core/src/write.jl
+++ b/core/src/write.jl
@@ -133,44 +133,26 @@ function basin_table(
 
     inflow_rate = FlatVector(saved.flow.saveval, :inflow)
     outflow_rate = FlatVector(saved.flow.saveval, :outflow)
-    precipitation = zeros(nrows)
-    evaporation = zeros(nrows)
-    drainage = zeros(nrows)
+    drainage = FlatVector(saved.flow.saveval, :drainage)
     infiltration = zeros(nrows)
-    balance_error = zeros(nrows)
-    relative_error = zeros(nrows)
+    evaporation = zeros(nrows)
+    precipitation = FlatVector(saved.flow.saveval, :precipitation)
+    storage_rate = FlatVector(saved.flow.saveval, :storage_rate)
+    balance_error = FlatVector(saved.flow.saveval, :balance_error)
+    relative_error = FlatVector(saved.flow.saveval, :relative_error)
 
     idx_row = 0
     for cvec in saved.flow.saveval
-        for (precipitation_, evaporation_, drainage_, infiltration_) in zip(
-            cvec.precipitation,
-            cvec.flow.evaporation,
-            cvec.drainage,
-            cvec.flow.infiltration,
-        )
+        for (evaporation_, infiltration_) in
+            zip(cvec.flow.evaporation, cvec.flow.infiltration)
             idx_row += 1
-            precipitation[idx_row] = precipitation_
             evaporation[idx_row] = evaporation_
-            drainage[idx_row] = drainage_
             infiltration[idx_row] = infiltration_
         end
     end
 
     time = repeat(data.time[begin:(end - 1)]; inner = nbasin)
-    Δtime_seconds = seconds.(diff(data.time))
-    Δtime = repeat(Δtime_seconds; inner = nbasin)
     node_id = repeat(Int32.(data.node_id); outer = ntsteps)
-    storage_rate = Δstorage ./ Δtime
-
-    for i in 1:nrows
-        total_in = inflow_rate[i] + precipitation[i] + drainage[i]
-        total_out = outflow_rate[i] + evaporation[i] + infiltration[i]
-        balance_error[i] = storage_rate[i] - (total_in - total_out)
-        mean_flow_rate = 0.5 * (total_in + total_out)
-        if mean_flow_rate != 0
-            relative_error[i] = balance_error[i] / mean_flow_rate
-        end
-    end
 
     return (;
         time,
diff --git a/core/test/run_models_test.jl b/core/test/run_models_test.jl
index 20cc0b492..f033b9f29 100644
--- a/core/test/run_models_test.jl
+++ b/core/test/run_models_test.jl
@@ -205,7 +205,7 @@ end
     @test successful_retcode(model)
     @test length(model.integrator.sol) == 2 # start and end
     @test model.integrator.p.basin.current_storage[Float64[]] ≈
-          Float32[803.7093, 803.68274, 495.241, 1318.3053] skip = Sys.isapple() atol = 1.5
+          Float32[828.5386, 801.88289, 492.290, 1318.3053] skip = Sys.isapple() atol = 1.5
 
     @test length(logger.logs) > 10
     @test logger.logs[1].level == Debug
@@ -242,7 +242,7 @@ end
     precipitation = model.integrator.p.basin.vertical_flux.precipitation
     @test length(precipitation) == 4
     @test model.integrator.p.basin.current_storage[parent(du)] ≈
-          Float32[697.30591, 697.2799, 419.19034, 1334.3859] atol = 2.0 skip = Sys.isapple()
+          Float32[720.23611, 694.8785, 415.22371, 1334.3859] atol = 2.0 skip = Sys.isapple()
 end
 
 @testitem "Allocation example model" begin
diff --git a/core/test/utils_test.jl b/core/test/utils_test.jl
index ef0b85ae8..8f5fa79f1 100644
--- a/core/test/utils_test.jl
+++ b/core/test/utils_test.jl
@@ -174,6 +174,7 @@ end
     t0 = 0.0
     u0 = Ribasim.build_state_vector(p)
     du0 = copy(u0)
+    p = Ribasim.build_flow_to_storage(p, u0)
     jac_prototype = Ribasim.get_jac_prototype(du0, u0, p, t0)
 
     # rows, cols, _ = findnz(jac_prototype)
@@ -195,6 +196,7 @@ end
     close(db)
     u0 = Ribasim.build_state_vector(p)
     du0 = copy(u0)
+    p = Ribasim.build_flow_to_storage(p, u0)
     jac_prototype = Ribasim.get_jac_prototype(du0, u0, p, t0)
 
     #! format: off
@@ -324,3 +326,39 @@ end
         end
     end
 end
+
+@testitem "flow_to_storage matrix" begin
+    using ComponentArrays: ComponentArray, Axis, getaxes
+    using LinearAlgebra: I
+    toml_path = normpath(@__DIR__, "../../generated_testmodels/basic/ribasim.toml")
+    @test ispath(toml_path)
+    model = Ribasim.Model(toml_path)
+    (; u, p) = model.integrator
+    n_basins = length(u.evaporation)
+    (; flow_to_storage) = p
+    flow_to_storage =
+        ComponentArray(flow_to_storage, (Axis(; basins = 1:n_basins), only(getaxes(u))))
+
+    @test flow_to_storage[:, :evaporation] == -I
+    @test flow_to_storage[:, :infiltration] == -I
+
+    for node_name in
+        [:tabulated_rating_curve, :pump, :outlet, :linear_resistance, :manning_resistance]
+        flow_to_storage_node = flow_to_storage[:, node_name]
+        # In every column there is either 0 or 1 instance of 1.0 (flow into a basin)
+        @test all(
+            i -> i ∈ (0, 1),
+            count(==(1.0), collect(flow_to_storage[:, :tabulated_rating_curve]); dims = 1),
+        )
+
+        # In every column there is either 0 or 1 instance of -1.0 (flow out of a basin)
+        @test all(
+            i -> i ∈ (0, 1),
+            count(
+                ==(1 - 0.0),
+                collect(flow_to_storage[:, :tabulated_rating_curve]);
+                dims = 1,
+            ),
+        )
+    end
+end
diff --git a/docs/reference/node/manning-resistance.qmd b/docs/reference/node/manning-resistance.qmd
index 63e66d979..94713b013 100644
--- a/docs/reference/node/manning-resistance.qmd
+++ b/docs/reference/node/manning-resistance.qmd
@@ -136,9 +136,10 @@ where $s$ is a relaxed square root function:
 
 $$
 s(x; x_0)
+=
 \begin{align}
   \begin{cases}
-    \frac{x}{4\sqrt{x_0}}\left(5-\left(\frac{x}{x_0}\right)^2\right) &\text{ if } |x| < x_0 \\
+    \frac{\sqrt{x_0}}{4}\left(\frac{x}{p}\right)^3\left(9 - 5\left(\frac{x}{p}\right)^2\right) &\text{ if } |x| < x_0 \\
     \textrm{sign}(x)\sqrt{|x|} &\text{ if } |x| \ge x_0
   \end{cases}
 \end{align}
@@ -151,7 +152,8 @@ import matplotlib.pyplot as plt
 
 def s(x, threshold):
   if np.abs(x) < threshold:
-    return 1/4 * (x / np.sqrt(threshold)) * (5.0 - (x / threshold)**2)
+    x_scaled = x / threshold
+    return np.sqrt(threshold) * x_scaled**3 * (9 - 5*x_scaled**2) / 4
   else:
     return np.sign(x)*np.sqrt(np.abs(x))
 
@@ -165,7 +167,7 @@ y_s = [s(x_, threshold) for x_ in x]
 
 ax.plot(x, y_o, ls = ":", label = r"sign$(x)\sqrt{|x|}$")
 ax.plot(x, y_s, color = "C0", label = r"$s\left(x; 10^{-3}\right)$")
-ax.legend()
+ax.legend();
 ```
 
 :::{.callout-note}