From 5981a1816abde80879963b0ac6908145fba0c80f Mon Sep 17 00:00:00 2001
From: Bart de Koning <74617371+SouthEndMusic@users.noreply.github.com>
Date: Wed, 21 Feb 2024 12:51:28 +0100
Subject: [PATCH] Basin allocation output (#1148)

Fixes https://github.com/Deltares/Ribasim/issues/1122.

# Upgrade information

The flow rate in `flow.arrow` is now in a column named `flow_rate`
rather than `flow` to be consistent with the input.
---
 core/src/allocation_optim.jl                  | 146 +++++---
 core/src/parameter.jl                         |  39 +-
 core/src/read.jl                              |  90 ++---
 core/src/util.jl                              |  30 +-
 core/src/write.jl                             |  70 ++--
 core/test/allocation_test.jl                  |  81 ++--
 core/test/run_models_test.jl                  |  37 +-
 core/test/time_test.jl                        |   4 +-
 core/test/utils_test.jl                       |   4 +
 core/test/validation_test.jl                  |   1 -
 docs/core/usage.qmd                           |  30 +-
 docs/python/examples.ipynb                    | 350 +++++++++++++++++-
 .../ribasim_testmodels/allocation.py          |   4 +-
 13 files changed, 690 insertions(+), 196 deletions(-)

diff --git a/core/src/allocation_optim.jl b/core/src/allocation_optim.jl
index 9e51d441b..a44c42146 100644
--- a/core/src/allocation_optim.jl
+++ b/core/src/allocation_optim.jl
@@ -126,12 +126,11 @@ function set_objective_priority!(
     priority_idx::Int,
 )::Nothing
     (; objective_type, problem, allocation_network_id) = allocation_model
-    (; graph, user, allocation) = p
+    (; graph, user, allocation, basin) = p
     (; demand_itp, demand_from_timeseries, node_id) = user
-    (; main_network_connections, subnetwork_demands, priorities) = allocation
+    (; main_network_connections, subnetwork_demands) = allocation
     edge_ids = graph[].edge_ids[allocation_network_id]
 
-    F = problem[:F]
     if objective_type in [:quadratic_absolute, :quadratic_relative]
         ex = JuMP.QuadExpr()
     elseif objective_type in [:linear_absolute, :linear_relative]
@@ -174,11 +173,12 @@ function set_objective_priority!(
     # Terms for basins
     F_basin_in = problem[:F_basin_in]
     for node_id in only(F_basin_in.axes)
-        priority_basin = get_basin_priority(p, node_id)
-        priority_now = priorities[priority_idx]
+        basin_priority_idx = get_basin_priority_idx(p, node_id)
         d =
-            priority_basin == priority_now ?
+            basin_priority_idx == priority_idx ?
             get_basin_demand(allocation_model, u, p, t, node_id) : 0.0
+        _, basin_idx = id_index(basin.node_id, node_id)
+        basin.demand[basin_idx] = d
         add_basin_term!(ex, problem, d, objective_type, node_id)
     end
 
@@ -193,7 +193,6 @@ Assign the allocations to the users as determined by the solution of the allocat
 function assign_allocations!(
     allocation_model::AllocationModel,
     p::Parameters,
-    t::Float64,
     priority_idx::Int;
     collect_demands::Bool = false,
 )::Nothing
@@ -205,7 +204,6 @@ function assign_allocations!(
         allocation_network_ids,
         main_network_connections,
     ) = allocation
-    (; record) = user
     edge_ids = graph[].edge_ids[allocation_network_id]
     main_network_source_edges = get_main_network_connections(p, allocation_network_id)
     F = problem[:F]
@@ -225,20 +223,6 @@ function assign_allocations!(
             allocated = JuMP.value(F[edge_id])
             user_idx = findsorted(user.node_id, user_node_id)
             user.allocated[user_idx][priority_idx] = allocated
-
-            # Save allocations to record
-            push!(record.time, t)
-            push!(record.subnetwork_id, allocation_model.allocation_network_id)
-            push!(record.user_node_id, Int(user_node_id))
-            push!(record.priority, allocation.priorities[priority_idx])
-            push!(record.demand, user.demand[user_idx])
-            push!(record.allocated, allocated)
-            # TODO: This is now the last abstraction before the allocation update,
-            # should be the average abstraction since the last allocation solve
-            push!(
-                record.abstracted,
-                get_flow(graph, inflow_id(graph, user_node_id), user_node_id, 0),
-            )
         end
     end
 
@@ -475,7 +459,76 @@ function adjust_basin_capacities!(
 end
 
 """
-Save the allocation flows per basin physical edge.
+Save the demands and allocated flows for users and basins.
+Note: Basin supply (negative demand) is only saved for the first priority.
+"""
+function save_demands_and_allocations!(
+    p::Parameters,
+    allocation_model::AllocationModel,
+    t::Float64,
+    priority_idx::Int,
+)::Nothing
+    (; graph, allocation, user, basin) = p
+    (; record_demand, priorities) = allocation
+    (; allocation_network_id, problem) = allocation_model
+    node_ids = graph[].node_ids[allocation_network_id]
+    constraints_outflow = problem[:basin_outflow]
+    F_basin_in = problem[:F_basin_in]
+    F_basin_out = problem[:F_basin_out]
+
+    for node_id in node_ids
+        has_demand = false
+
+        if node_id.type == NodeType.User
+            has_demand = true
+            user_idx = findsorted(user.node_id, node_id)
+            demand = user.demand[user_idx]
+            allocated = user.allocated[user_idx][priority_idx]
+            realized = get_flow(graph, inflow_id(graph, node_id), node_id, 0)
+
+        elseif node_id.type == NodeType.Basin
+            basin_priority_idx = get_basin_priority_idx(p, node_id)
+
+            if priority_idx == 1 || basin_priority_idx == priority_idx
+                has_demand = true
+                demand = 0.0
+                if priority_idx == 1
+                    # Basin surplus
+                    demand -= JuMP.normalized_rhs(constraints_outflow[node_id])
+                end
+                if priority_idx == basin_priority_idx
+                    # Basin demand
+                    _, basin_idx = id_index(basin.node_id, node_id)
+                    demand += basin.demand[basin_idx]
+                end
+                allocated =
+                    JuMP.value(F_basin_in[node_id]) - JuMP.value(F_basin_out[node_id])
+                # TODO: realized for a basin is not so clear, maybe it should be Δstorage/Δt
+                # over the last allocation interval?
+                realized = 0.0
+            end
+        end
+
+        if has_demand
+            # Save allocations and demands to record
+            push!(record_demand.time, t)
+            push!(record_demand.subnetwork_id, allocation_network_id)
+            push!(record_demand.node_type, string(node_id.type))
+            push!(record_demand.node_id, Int(node_id))
+            push!(record_demand.priority, priorities[priority_idx])
+            push!(record_demand.demand, demand)
+            push!(record_demand.allocated, allocated)
+
+            # TODO: This is now the last abstraction before the allocation update,
+            # should be the average abstraction since the last allocation solve
+            push!(record_demand.realized, realized)
+        end
+    end
+    return nothing
+end
+
+"""
+Save the allocation flows per basin and physical edge.
 """
 function save_allocation_flows!(
     p::Parameters,
@@ -486,41 +539,45 @@ function save_allocation_flows!(
 )::Nothing
     (; problem, allocation_network_id) = allocation_model
     (; allocation, graph) = p
-    (; record) = allocation
+    (; record_flow) = allocation
     F = problem[:F]
     F_basin_in = problem[:F_basin_in]
     F_basin_out = problem[:F_basin_out]
 
     # Edge flows
     for allocation_edge in first(F.axes)
-        flow = JuMP.value(F[allocation_edge])
+        flow_rate = JuMP.value(F[allocation_edge])
         edge_metadata = graph[allocation_edge...]
         (; node_ids) = edge_metadata
 
         for i in eachindex(node_ids)[1:(end - 1)]
-            push!(record.time, t)
-            push!(record.edge_id, edge_metadata.id)
-            push!(record.from_node_id, node_ids[i])
-            push!(record.to_node_id, node_ids[i + 1])
-            push!(record.subnetwork_id, allocation_network_id)
-            push!(record.priority, priority)
-            push!(record.flow, flow)
-            push!(record.collect_demands, collect_demands)
+            push!(record_flow.time, t)
+            push!(record_flow.edge_id, edge_metadata.id)
+            push!(record_flow.from_node_type, string(node_ids[i].type))
+            push!(record_flow.from_node_id, Int(node_ids[i]))
+            push!(record_flow.to_node_type, string(node_ids[i + 1].type))
+            push!(record_flow.to_node_id, Int(node_ids[i + 1]))
+            push!(record_flow.subnetwork_id, allocation_network_id)
+            push!(record_flow.priority, priority)
+            push!(record_flow.flow_rate, flow_rate)
+            push!(record_flow.collect_demands, collect_demands)
         end
     end
 
     # Basin flows
     for node_id in graph[].node_ids[allocation_network_id]
         if node_id.type == NodeType.Basin
-            flow = JuMP.value(F_basin_out[node_id]) - JuMP.value(F_basin_in[node_id])
-            push!(record.time, t)
-            push!(record.edge_id, 0)
-            push!(record.from_node_id, node_id)
-            push!(record.to_node_id, node_id)
-            push!(record.subnetwork_id, allocation_network_id)
-            push!(record.priority, priority)
-            push!(record.flow, flow)
-            push!(record.collect_demands, collect_demands)
+            flow_rate = JuMP.value(F_basin_out[node_id]) - JuMP.value(F_basin_in[node_id])
+            push!(record_flow.time, t)
+            push!(record_flow.edge_id, 0)
+            push!(record_flow.from_node_type, string(NodeType.Basin))
+            push!(record_flow.from_node_id, node_id)
+            push!(record_flow.to_node_type, string(NodeType.Basin))
+            push!(record_flow.to_node_id, node_id)
+            push!(record_flow.subnetwork_id, allocation_network_id)
+            push!(record_flow.priority, priority)
+            push!(record_flow.flow_rate, flow_rate)
+            push!(record_flow.collect_demands, collect_demands)
         end
     end
 
@@ -580,7 +637,10 @@ function allocate!(
         end
 
         # Assign the allocations to the users for this priority
-        assign_allocations!(allocation_model, p, t, priority_idx; collect_demands)
+        assign_allocations!(allocation_model, p, priority_idx; collect_demands)
+
+        # Save the demands and allocated flows for all nodes that have these
+        save_demands_and_allocations!(p, allocation_model, t, priority_idx)
 
         # Save the flows over all edges in the subnetwork
         save_allocation_flows!(
diff --git a/core/src/parameter.jl b/core/src/parameter.jl
index fc8c8feb7..b996fd72f 100644
--- a/core/src/parameter.jl
+++ b/core/src/parameter.jl
@@ -65,7 +65,8 @@ allocation models: The allocation models for the main network and subnetworks co
 main_network_connections: (from_id, to_id) from the main network to the subnetwork per subnetwork
 priorities: All used priority values.
 subnetwork_demands: The demand of an edge from the main network to a subnetwork
-record: A record of all flows computed by allocation optimization, eventually saved to
+record_demand: A record of demands and allocated flows for nodes that have these.
+record_flow: A record of all flows computed by allocation optimization, eventually saved to
     output file
 """
 struct Allocation
@@ -75,14 +76,26 @@ struct Allocation
     priorities::Vector{Int}
     subnetwork_demands::Dict{Tuple{NodeID, NodeID}, Vector{Float64}}
     subnetwork_allocateds::Dict{Tuple{NodeID, NodeID}, Vector{Float64}}
-    record::@NamedTuple{
+    record_demand::@NamedTuple{
+        time::Vector{Float64},
+        subnetwork_id::Vector{Int},
+        node_type::Vector{String},
+        node_id::Vector{Int},
+        priority::Vector{Int},
+        demand::Vector{Float64},
+        allocated::Vector{Float64},
+        realized::Vector{Float64},
+    }
+    record_flow::@NamedTuple{
         time::Vector{Float64},
         edge_id::Vector{Int},
+        from_node_type::Vector{String},
         from_node_id::Vector{Int},
+        to_node_type::Vector{String},
         to_node_id::Vector{Int},
         subnetwork_id::Vector{Int},
         priority::Vector{Int},
-        flow::Vector{Float64},
+        flow_rate::Vector{Float64},
         collect_demands::BitVector,
     }
 end
@@ -146,14 +159,16 @@ struct Basin{T, C} <: AbstractParameterNode
     potential_evaporation::Vector{Float64}
     drainage::Vector{Float64}
     infiltration::Vector{Float64}
-    # cache this to avoid recomputation
+    # Cache this to avoid recomputation
     current_level::T
     current_area::T
     # Discrete values for interpolation
     area::Vector{Vector{Float64}}
     level::Vector{Vector{Float64}}
     storage::Vector{Vector{Float64}}
-    # data source for parameter updates
+    # Demands and allocated flows for allocation if applicable
+    demand::Vector{Float64}
+    # Data source for parameter updates
     time::StructVector{BasinTimeV1, C, Int}
 
     function Basin(
@@ -167,6 +182,7 @@ struct Basin{T, C} <: AbstractParameterNode
         area,
         level,
         storage,
+        demand,
         time::StructVector{BasinTimeV1, C, Int},
     ) where {T, C}
         is_valid = valid_profiles(node_id, level, area)
@@ -182,6 +198,7 @@ struct Basin{T, C} <: AbstractParameterNode
             area,
             level,
             storage,
+            demand,
             time,
         )
     end
@@ -467,7 +484,6 @@ active: whether this node is active and thus demands water
 allocated: water flux currently allocated to user per priority
 return_factor: the factor in [0,1] of how much of the abstracted water is given back to the system
 min_level: The level of the source basin below which the user does not abstract
-record: Collected data of allocation optimizations for output file.
 """
 struct User <: AbstractParameterNode
     node_id::Vector{NodeID}
@@ -478,15 +494,6 @@ struct User <: AbstractParameterNode
     allocated::Vector{Vector{Float64}}
     return_factor::Vector{Float64}
     min_level::Vector{Float64}
-    record::@NamedTuple{
-        time::Vector{Float64},
-        subnetwork_id::Vector{Int},
-        user_node_id::Vector{Int},
-        priority::Vector{Int},
-        demand::Vector{Float64},
-        allocated::Vector{Float64},
-        abstracted::Vector{Float64},
-    }
 
     function User(
         node_id,
@@ -498,7 +505,6 @@ struct User <: AbstractParameterNode
         return_factor,
         min_level,
         priorities,
-        record,
     )
         if valid_demand(node_id, demand_itp, priorities)
             return new(
@@ -510,7 +516,6 @@ struct User <: AbstractParameterNode
                 allocated,
                 return_factor,
                 min_level,
-                record,
             )
         else
             error("Invalid demand")
diff --git a/core/src/read.jl b/core/src/read.jl
index c6df7eb85..c569a617f 100644
--- a/core/src/read.jl
+++ b/core/src/read.jl
@@ -513,6 +513,8 @@ function Basin(db::DB, config::Config, chunk_sizes::Vector{Int})::Basin
     set_current_value!(table, node_id, time, config.starttime)
     check_no_nans(table, "Basin")
 
+    demand = zeros(length(node_id))
+
     return Basin(
         Indices(NodeID.(NodeType.Basin, node_id)),
         precipitation,
@@ -524,6 +526,7 @@ function Basin(db::DB, config::Config, chunk_sizes::Vector{Int})::Basin
         area,
         level,
         storage,
+        demand,
         time,
     )
 end
@@ -741,16 +744,6 @@ function User(db::DB, config::Config)::User
 
     allocated = [fill(Inf, length(priorities)) for id in node_ids]
 
-    record = (
-        time = Float64[],
-        subnetwork_id = Int[],
-        user_node_id = Int[],
-        priority = Int[],
-        demand = Float64[],
-        allocated = Float64[],
-        abstracted = Float64[],
-    )
-
     return User(
         node_ids,
         active,
@@ -761,7 +754,6 @@ function User(db::DB, config::Config)::User
         return_factor,
         min_level,
         priorities,
-        record,
     )
 end
 
@@ -823,6 +815,45 @@ function Subgrid(db::DB, config::Config, basin::Basin)::Subgrid
     return Subgrid(basin_ids, interpolations, fill(NaN, length(basin_ids)))
 end
 
+function Allocation(db::DB, config::Config)::Allocation
+    record_demand = (
+        time = Float64[],
+        subnetwork_id = Int[],
+        node_type = String[],
+        node_id = Int[],
+        priority = Int[],
+        demand = Float64[],
+        allocated = Float64[],
+        realized = Float64[],
+    )
+
+    record_flow = (
+        time = Float64[],
+        edge_id = Int[],
+        from_node_type = String[],
+        from_node_id = Int[],
+        to_node_type = String[],
+        to_node_id = Int[],
+        subnetwork_id = Int[],
+        priority = Int[],
+        flow_rate = Float64[],
+        collect_demands = BitVector(),
+    )
+
+    allocation = Allocation(
+        Int[],
+        AllocationModel[],
+        Vector{Tuple{NodeID, NodeID}}[],
+        get_all_priorities(db, config),
+        Dict{Tuple{NodeID, NodeID}, Float64}(),
+        Dict{Tuple{NodeID, NodeID}, Float64}(),
+        record_demand,
+        record_flow,
+    )
+
+    return allocation
+end
+
 """
 Get the chunk sizes for DiffCache; differentiation w.r.t. u
 and t (the latter only if a Rosenbrock algorithm is used).
@@ -840,24 +871,7 @@ function Parameters(db::DB, config::Config)::Parameters
     n_states = length(get_ids(db, "Basin")) + length(get_ids(db, "PidControl"))
     chunk_sizes = get_chunk_sizes(config, n_states)
     graph = create_graph(db, config, chunk_sizes)
-    allocation = Allocation(
-        Int[],
-        AllocationModel[],
-        Vector{Tuple{NodeID, NodeID}}[],
-        get_all_priorities(db, config),
-        Dict{Tuple{NodeID, NodeID}, Float64}(),
-        Dict{Tuple{NodeID, NodeID}, Float64}(),
-        (;
-            time = Float64[],
-            edge_id = Int[],
-            from_node_id = Int[],
-            to_node_id = Int[],
-            subnetwork_id = Int[],
-            priority = Int[],
-            flow = Float64[],
-            collect_demands = BitVector(),
-        ),
-    )
+    allocation = Allocation(db, config)
 
     if !valid_edges(graph)
         error("Invalid edge(s) found.")
@@ -880,22 +894,6 @@ function Parameters(db::DB, config::Config)::Parameters
     basin = Basin(db, config, chunk_sizes)
     subgrid_level = Subgrid(db, config, basin)
 
-    # Set is_pid_controlled to true for those pumps and outlets that are PID controlled
-    for id in pid_control.node_id
-        id_controlled = only(outneighbor_labels_type(graph, id, EdgeType.control))
-        if id_controlled.type == NodeType.Pump
-            pump_idx = findsorted(pump.node_id, id_controlled)
-            pump.is_pid_controlled[pump_idx] = true
-        elseif id_controlled.type == NodeType.Outlet
-            outlet_idx = findsorted(outlet.node_id, id_controlled)
-            outlet.is_pid_controlled[outlet_idx] = true
-        else
-            error(
-                "Only Pump and Outlet can be controlled by PidController, got $is_controlled",
-            )
-        end
-    end
-
     p = Parameters(
         config.starttime,
         graph,
@@ -917,6 +915,8 @@ function Parameters(db::DB, config::Config)::Parameters
         subgrid_level,
     )
 
+    set_is_pid_controlled!(p)
+
     if !valid_n_neighbors(p)
         error("Invalid number of connections for certain node types.")
     end
diff --git a/core/src/util.jl b/core/src/util.jl
index e9165203e..4deafd45e 100644
--- a/core/src/util.jl
+++ b/core/src/util.jl
@@ -652,14 +652,38 @@ function get_all_priorities(db::DB, config::Config)::Vector{Int}
     return sort(unique(priorities))
 end
 
-function get_basin_priority(p::Parameters, node_id::NodeID)::Int
-    (; graph, target_level) = p
+function get_basin_priority_idx(p::Parameters, node_id::NodeID)::Int
+    (; graph, target_level, allocation) = p
     @assert node_id.type == NodeType.Basin
     inneighbors_control = inneighbor_labels_type(graph, node_id, EdgeType.control)
     if isempty(inneighbors_control)
         return 0
     else
         idx = findsorted(target_level.node_id, only(inneighbors_control))
-        return target_level.priority[idx]
+        priority = target_level.priority[idx]
+        return findsorted(allocation.priorities, priority)
     end
 end
+
+"""
+Set is_pid_controlled to true for those pumps and outlets that are PID controlled
+"""
+function set_is_pid_controlled!(p::Parameters)::Nothing
+    (; graph, pid_control, pump, outlet) = p
+
+    for id in pid_control.node_id
+        id_controlled = only(outneighbor_labels_type(graph, id, EdgeType.control))
+        if id_controlled.type == NodeType.Pump
+            pump_idx = findsorted(pump.node_id, id_controlled)
+            pump.is_pid_controlled[pump_idx] = true
+        elseif id_controlled.type == NodeType.Outlet
+            outlet_idx = findsorted(outlet.node_id, id_controlled)
+            outlet.is_pid_controlled[outlet_idx] = true
+        else
+            error(
+                "Only Pump and Outlet can be controlled by PidController, got $is_controlled",
+            )
+        end
+    end
+    return nothing
+end
diff --git a/core/src/write.jl b/core/src/write.jl
index a4b319b7f..17f21a8bf 100644
--- a/core/src/write.jl
+++ b/core/src/write.jl
@@ -101,9 +101,11 @@ function flow_table(
 )::@NamedTuple{
     time::Vector{DateTime},
     edge_id::Vector{Union{Int, Missing}},
+    from_node_type::Vector{String},
     from_node_id::Vector{Int},
+    to_node_type::Vector{String},
     to_node_id::Vector{Int},
-    flow::FlatVector{Float64},
+    flow_rate::FlatVector{Float64},
 }
     (; config, saved, integrator) = model
     (; t, saveval) = saved.flow
@@ -111,7 +113,9 @@ function flow_table(
     (; flow_dict, flow_vertical_dict) = graph[]
 
     # self-loops have no edge ID
+    from_node_type = String[]
     from_node_id = Int[]
+    to_node_type = String[]
     to_node_id = Int[]
     unique_edge_ids_flow = Union{Int, Missing}[]
 
@@ -121,7 +125,9 @@ function flow_table(
     end
 
     for id in vertical_flow_node_ids
+        push!(from_node_type, string(id.type))
         push!(from_node_id, id.value)
+        push!(to_node_type, string(id.type))
         push!(to_node_id, id.value)
         push!(unique_edge_ids_flow, missing)
     end
@@ -132,7 +138,9 @@ function flow_table(
     end
 
     for (from_id, to_id) in flow_edge_ids
+        push!(from_node_type, string(from_id.type))
         push!(from_node_id, from_id.value)
+        push!(to_node_type, string(to_id.type))
         push!(to_node_id, to_id.value)
         push!(unique_edge_ids_flow, graph[from_id, to_id].id)
     end
@@ -142,11 +150,21 @@ function flow_table(
 
     time = repeat(datetime_since.(t, config.starttime); inner = nflow)
     edge_id = repeat(unique_edge_ids_flow; outer = ntsteps)
+    from_node_type = repeat(from_node_type; outer = ntsteps)
     from_node_id = repeat(from_node_id; outer = ntsteps)
+    to_node_type = repeat(to_node_type; outer = ntsteps)
     to_node_id = repeat(to_node_id; outer = ntsteps)
-    flow = FlatVector(saveval)
+    flow_rate = FlatVector(saveval)
 
-    return (; time, edge_id, from_node_id, to_node_id, flow)
+    return (;
+        time,
+        edge_id,
+        from_node_type,
+        from_node_id,
+        to_node_type,
+        to_node_id,
+        flow_rate,
+    )
 end
 
 "Create a discrete control result table from the saved data"
@@ -171,24 +189,26 @@ function allocation_table(
 )::@NamedTuple{
     time::Vector{DateTime},
     subnetwork_id::Vector{Int},
-    user_node_id::Vector{Int},
+    node_type::Vector{String},
+    node_id::Vector{Int},
     priority::Vector{Int},
     demand::Vector{Float64},
     allocated::Vector{Float64},
-    abstracted::Vector{Float64},
+    realized::Vector{Float64},
 }
     (; config) = model
-    (; record) = model.integrator.p.user
+    (; record_demand) = model.integrator.p.allocation
 
-    time = datetime_since.(record.time, config.starttime)
+    time = datetime_since.(record_demand.time, config.starttime)
     return (;
         time,
-        record.subnetwork_id,
-        record.user_node_id,
-        record.priority,
-        record.demand,
-        record.allocated,
-        record.abstracted,
+        record_demand.subnetwork_id,
+        record_demand.node_type,
+        record_demand.node_id,
+        record_demand.priority,
+        record_demand.demand,
+        record_demand.allocated,
+        record_demand.realized,
     )
 end
 
@@ -197,27 +217,31 @@ function allocation_flow_table(
 )::@NamedTuple{
     time::Vector{DateTime},
     edge_id::Vector{Int},
+    from_node_type::Vector{String},
     from_node_id::Vector{Int},
+    to_node_type::Vector{String},
     to_node_id::Vector{Int},
     subnetwork_id::Vector{Int},
     priority::Vector{Int},
-    flow::Vector{Float64},
+    flow_rate::Vector{Float64},
     collect_demands::BitVector,
 }
     (; config) = model
-    (; record) = model.integrator.p.allocation
+    (; record_flow) = model.integrator.p.allocation
 
-    time = datetime_since.(record.time, config.starttime)
+    time = datetime_since.(record_flow.time, config.starttime)
 
     return (;
         time,
-        record.edge_id,
-        record.from_node_id,
-        record.to_node_id,
-        record.subnetwork_id,
-        record.priority,
-        record.flow,
-        record.collect_demands,
+        record_flow.edge_id,
+        record_flow.from_node_type,
+        record_flow.from_node_id,
+        record_flow.to_node_type,
+        record_flow.to_node_id,
+        record_flow.subnetwork_id,
+        record_flow.priority,
+        record_flow.flow_rate,
+        record_flow.collect_demands,
     )
 end
 
diff --git a/core/test/allocation_test.jl b/core/test/allocation_test.jl
index cf12a35ca..9aa5e74f6 100644
--- a/core/test/allocation_test.jl
+++ b/core/test/allocation_test.jl
@@ -175,17 +175,21 @@ end
     ) ≈ -0.25
 
     solve!(model)
-    record_allocation = DataFrame(model.integrator.p.user.record)
+    record_allocation = DataFrame(model.integrator.p.allocation.record_demand)
     record_control = model.integrator.p.discrete_control.record
-    groups = groupby(record_allocation, [:user_node_id, :priority])
+    groups = groupby(record_allocation, [:node_type, :node_id, :priority])
     fractional_flow = model.integrator.p.fractional_flow
     (; control_mapping) = fractional_flow
     t_control = record_control.time[2]
 
-    allocated_6_before = groups[(6, 1)][groups[(6, 1)].time .< t_control, :].allocated
-    allocated_9_before = groups[(9, 1)][groups[(9, 1)].time .< t_control, :].allocated
-    allocated_6_after = groups[(6, 1)][groups[(6, 1)].time .> t_control, :].allocated
-    allocated_9_after = groups[(9, 1)][groups[(9, 1)].time .> t_control, :].allocated
+    allocated_6_before =
+        groups[("User", 6, 1)][groups[("User", 6, 1)].time .< t_control, :].allocated
+    allocated_9_before =
+        groups[("User", 9, 1)][groups[("User", 9, 1)].time .< t_control, :].allocated
+    allocated_6_after =
+        groups[("User", 6, 1)][groups[("User", 6, 1)].time .> t_control, :].allocated
+    allocated_9_after =
+        groups[("User", 9, 1)][groups[("User", 9, 1)].time .> t_control, :].allocated
     @test all(
         allocated_9_before ./ allocated_6_before .<=
         control_mapping[(NodeID(:FractionalFlow, 7), "A")].fraction /
@@ -344,28 +348,45 @@ end
     Δt_allocation = allocation.allocation_models[1].Δt_allocation
 
     # Until the first allocation solve, the user abstracts fully
-    pre_allocation = t .<= Δt_allocation
-    u_pre_allocation(τ) = storage[1] + (q + ϕ - d) * τ
-    @test storage[pre_allocation] ≈ u_pre_allocation.(t[pre_allocation]) rtol = 1e-4
-
-    # Until the basin is at its maximum level, the user does not abstract
-    basin_filling = @. ~pre_allocation && (storage <= A * l_max)
-    fill_start_idx = findlast(pre_allocation)
-    u_filling(τ) = storage[fill_start_idx] + (q + ϕ) * (τ - t[fill_start_idx])
-    @test storage[basin_filling] ≈ u_filling.(t[basin_filling]) rtol = 1e-4
-
-    # After the basin has reached its maximum level, the user abstracts fully again
-    precipitation = eachindex(storage) .<= argmax(storage)
-    after_filling = @. ~pre_allocation && ~basin_filling && precipitation
-    fill_stop_idx = findfirst(after_filling)
-    u_after_filling(τ) = storage[fill_stop_idx] + (q + ϕ - d) * (τ - t[fill_stop_idx])
-    @test storage[after_filling] ≈ u_after_filling.(t[after_filling]) rtol = 1e-4
-
-    # After precipitation stops, the user still abstracts from the basin so the storage decreases
-    storage_reduction = @. ~precipitation && t <= 1.8e6
-    storage_reduction_start = findfirst(storage_reduction)
-    u_storage_reduction(τ) =
-        storage[storage_reduction_start] + (q - d) * (τ - t[storage_reduction_start])
-    @test storage[storage_reduction] ≈ u_storage_reduction.(t[storage_reduction]) rtol =
-        1e-4
+    stage_1 = t .<= Δt_allocation
+    u_stage_1(τ) = storage[1] + (q + ϕ - d) * τ
+    @test storage[stage_1] ≈ u_stage_1.(t[stage_1]) rtol = 1e-4
+
+    # In this section the basin leaves no supply for the user
+    stage_2 = Δt_allocation .<= t .<= 3 * Δt_allocation
+    stage_2_start_idx = findfirst(stage_2)
+    u_stage_2(τ) = storage[stage_2_start_idx] + (q + ϕ) * (τ - t[stage_2_start_idx])
+    @test storage[stage_2] ≈ u_stage_2.(t[stage_2]) rtol = 1e-4
+
+    # In this section (and following sections) the basin has no longer a (positive) demand,
+    # since precipitation provides enough water to get the basin to its target level
+    # The FlowBoundary flow gets fully allocated to the user
+    stage_3 = 3 * Δt_allocation .<= t .<= 8 * Δt_allocation
+    stage_3_start_idx = findfirst(stage_3)
+    u_stage_3(τ) = storage[stage_3_start_idx] + ϕ * (τ - t[stage_3_start_idx])
+    @test storage[stage_3] ≈ u_stage_3.(t[stage_3]) rtol = 1e-4
+
+    # In this section the basin enters its surplus stage,
+    # even though initially the level is below the maximum level. This is because the simulation
+    # anticipates that the current precipitation is going to bring the basin level over
+    # its maximum level
+    stage_4 = 8 * Δt_allocation .<= t .<= 12 * Δt_allocation
+    stage_4_start_idx = findfirst(stage_4)
+    u_stage_4(τ) = storage[stage_4_start_idx] + (q + ϕ - d) * (τ - t[stage_4_start_idx])
+    @test storage[stage_4] ≈ u_stage_4.(t[stage_4]) rtol = 1e-4
+
+    # At the start of this section precipitation stops, and so the user
+    # partly uses surplus water from the basin to fulfill its demand
+    stage_5 = 13 * Δt_allocation .<= t .<= 16 * Δt_allocation
+    stage_5_start_idx = findfirst(stage_5)
+    u_stage_5(τ) = storage[stage_5_start_idx] + (q - d) * (τ - t[stage_5_start_idx])
+    @test storage[stage_5] ≈ u_stage_5.(t[stage_5]) rtol = 1e-4
+
+    # From this point the basin is in a dynamical equilibrium,
+    # since the basin has no supply so the user abstracts precisely
+    # the flow from the level boundary
+    stage_6 = 17 * Δt_allocation .<= t
+    stage_6_start_idx = findfirst(stage_6)
+    u_stage_6(τ) = storage[stage_6_start_idx]
+    @test storage[stage_6] ≈ u_stage_6.(t[stage_6]) rtol = 1e-4
 end
diff --git a/core/test/run_models_test.jl b/core/test/run_models_test.jl
index e08b152e5..051a27d75 100644
--- a/core/test/run_models_test.jl
+++ b/core/test/run_models_test.jl
@@ -32,8 +32,16 @@
 
     @testset "Schema" begin
         @test Tables.schema(flow) == Tables.Schema(
-            (:time, :edge_id, :from_node_id, :to_node_id, :flow),
-            (DateTime, Union{Int, Missing}, Int, Int, Float64),
+            (
+                :time,
+                :edge_id,
+                :from_node_type,
+                :from_node_id,
+                :to_node_type,
+                :to_node_id,
+                :flow_rate,
+            ),
+            (DateTime, Union{Int, Missing}, String, Int, String, Int, Float64),
         )
         @test Tables.schema(basin) == Tables.Schema(
             (:time, :node_id, :storage, :level),
@@ -47,26 +55,29 @@
             (
                 :time,
                 :subnetwork_id,
-                :user_node_id,
+                :node_type,
+                :node_id,
                 :priority,
                 :demand,
                 :allocated,
-                :abstracted,
+                :realized,
             ),
-            (DateTime, Int, Int, Int, Float64, Float64, Float64),
+            (DateTime, Int, String, Int, Int, Float64, Float64, Float64),
         )
         @test Tables.schema(allocation_flow) == Tables.Schema(
             (
                 :time,
                 :edge_id,
+                :from_node_type,
                 :from_node_id,
+                :to_node_type,
                 :to_node_id,
                 :subnetwork_id,
                 :priority,
-                :flow,
+                :flow_rate,
                 :collect_demands,
             ),
-            (DateTime, Int, Int, Int, Int, Int, Float64, Bool),
+            (DateTime, Int, String, Int, String, Int, Int, Int, Float64, Bool),
         )
         @test Tables.schema(subgrid) ==
               Tables.Schema((:time, :subgrid_id, :subgrid_level), (DateTime, Int, Float64))
@@ -187,10 +198,7 @@ end
     @test logger.logs[1].message == "Read database into memory."
 
     table = Ribasim.flow_table(model)
-    @test Tables.schema(table) == Tables.Schema(
-        (:time, :edge_id, :from_node_id, :to_node_id, :flow),
-        (DateTime, Union{Int, Missing}, Int, Int, Float64),
-    )
+
     # flows are recorded at the end of each period, and are undefined at the start
     @test unique(table.time) == Ribasim.datetimes(model)[2:end]
 
@@ -198,8 +206,9 @@ end
     t = table.time[1]
     @test length(p.fractional_flow.node_id) == 3
     for id in p.fractional_flow.node_id
-        inflow = only(table.flow[table.to_node_id .== id.value .&& table.time .== t])
-        outflow = only(table.flow[table.from_node_id .== id.value .&& table.time .== t])
+        inflow = only(table.flow_rate[table.to_node_id .== id.value .&& table.time .== t])
+        outflow =
+            only(table.flow_rate[table.from_node_id .== id.value .&& table.time .== t])
         @test inflow == outflow
     end
 end
@@ -355,7 +364,7 @@ end
         level.t[2] * (outlet.min_crest_level[1] - level.u[1]) / (level.u[2] - level.u[1])
 
     # No outlet flow when upstream level is below minimum crest level
-    @test all(@. outlet_flow.flow[timesteps <= t_min_crest_level] == 0)
+    @test all(@. outlet_flow.flow_rate[timesteps <= t_min_crest_level] == 0)
 
     timesteps = Ribasim.timesteps(model)
     t_maximum_level = level.t[2]
diff --git a/core/test/time_test.jl b/core/test/time_test.jl
index 02f73af5f..82ad216d3 100644
--- a/core/test/time_test.jl
+++ b/core/test/time_test.jl
@@ -18,13 +18,13 @@
             [:time, :from_node_id, :to_node_id] =>
                 (t, from, to) -> is_summer(t) && from === 1 && to === 1,
             flow,
-        ).flow
+        ).flow_rate
     flow_1_to_2 = filter(
         [:time, :from_node_id, :to_node_id] =>
             (t, from, to) -> is_summer(t) && from === 1 && to === 2,
         flow,
     )
-    @test flow_added_1 == flow_1_to_2.flow
+    @test flow_added_1 == flow_1_to_2.flow_rate
 
     t = Ribasim.seconds_since.(flow_1_to_2.time, model.config.starttime)
     flow_expected = @. 1 + sin(0.5 * π * (t - t[1]) / (t[end] - t[1]))^2
diff --git a/core/test/utils_test.jl b/core/test/utils_test.jl
index a9784bc20..6c23e7558 100644
--- a/core/test/utils_test.jl
+++ b/core/test/utils_test.jl
@@ -35,6 +35,7 @@ end
     level = [[0.0, 1.0], [4.0, 5.0]]
     darea = zeros(2)
     storage = Ribasim.profile_storage.(level, area)
+    demand = zeros(2)
     basin = Ribasim.Basin(
         Indices(NodeID.(:Basin, [5, 7])),
         [2.0, 3.0],
@@ -46,6 +47,7 @@ end
         area,
         level,
         storage,
+        demand,
         StructVector{Ribasim.BasinTimeV1}(undef, 0),
     )
 
@@ -110,6 +112,7 @@ end
         1.1,
     ]
     storage = Ribasim.profile_storage(level, area)
+    demand = zeros(1)
     basin = Ribasim.Basin(
         Indices(NodeID.(:Basin, [1])),
         zeros(1),
@@ -121,6 +124,7 @@ end
         [area],
         [level],
         [storage],
+        demand,
         StructVector{Ribasim.BasinTimeV1}(undef, 0),
     )
 
diff --git a/core/test/validation_test.jl b/core/test/validation_test.jl
index 2f7f5a250..aec09f671 100644
--- a/core/test/validation_test.jl
+++ b/core/test/validation_test.jl
@@ -421,7 +421,6 @@ end
             [0.9],
             [0.9],
             [1],
-            [],
         )
     end
 
diff --git a/docs/core/usage.qmd b/docs/core/usage.qmd
index 91250c4d8..c9a320ff2 100644
--- a/docs/core/usage.qmd
+++ b/docs/core/usage.qmd
@@ -684,13 +684,15 @@ The table is sorted by time, and per time it is sorted by `node_id`.
 The flow table contains results of the flow on every edge in the model, for each solver
 timestep.
 
-column        | type                | unit
-------------- | ------------------- | ----
-time          | DateTime            | -
-edge_id       | Union{Int, Missing} | -
-from_node_id  | Int                 | -
-to_node_id    | Int                 | -
-flow          | Float64             | $m^3 s^{-1}$
+column         | type                | unit
+-------------- | ------------------- | ----
+time           | DateTime            | -
+edge_id        | Union{Int, Missing} | -
+from_node_type | String              | -
+from_node_id   | Int                 | -
+to_node_type   | String              | -
+to_node_id     | Int                 | -
+flow_rate      | Float64             | $m^3 s^{-1}$
 
 The table is sorted by time, and per time the same `edge_id` order is used, though not sorted.
 Flows that are added to the model at a node, have a missing `edge_id`, and identical `from_node_id` and `to_node_id`.
@@ -709,17 +711,21 @@ control_state   | String
 
 ## Allocation - `allocation.arrow`
 
-The allocation table contains a record of allocation results: when it happened, for which user node, in which allocation network, and what the demand, allocated flow and abstracted flow were.
+The allocation table contains a record of allocation results: when it happened, for which node, in which allocation network, and what the demand, allocated flow and realized flow were.
 
 column        | type
 --------------| -----
 time          | DateTime
 subnetwork_id | Int
-user_node_id  | Int
+node_type     | String
+node_id       | Int
 priority      | Int
 demand        | Float64
 allocated     | Float64
-abstracted    | Float64
+realized      | Float64
+
+For Basins the values `demand`, `allocated` and `realized` are positive if the Basin level is below the minimum level given by a `TargetLevel` node.
+The values are negative if the Basin supplies due to a surplus of water.
 
 ::: {.callout-note}
 Currently the stored demand and abstraction rate are those at the allocation timepoint (and the abstraction rate is based on the previous allocation optimization). In the future these will be an average over the previous allocation timestep.
@@ -735,9 +741,11 @@ column          | type
 ----------------| -----
 time            | DateTime
 edge_id         | Int
+from_node_type  | String
 from_node_id    | Int
+to_node_type    | String
 to_node_id      | Int
 subnetwork_id   | Int
 priority        | Int
-flow            | Float64
+flow_rate       | Float64
 collect_demands | Bool
diff --git a/docs/python/examples.ipynb b/docs/python/examples.ipynb
index 691df7f59..da8f3403a 100644
--- a/docs/python/examples.ipynb
+++ b/docs/python/examples.ipynb
@@ -609,7 +609,7 @@
    "source": [
     "df_flow = pd.read_feather(datadir / \"basic_transient/results/flow.arrow\")\n",
     "df_flow[\"edge\"] = list(zip(df_flow.from_node_id, df_flow.to_node_id))\n",
-    "df_flow[\"flow_m3d\"] = df_flow.flow * 86400\n",
+    "df_flow[\"flow_m3d\"] = df_flow.flow_rate * 86400\n",
     "ax = df_flow.pivot_table(index=\"time\", columns=\"edge\", values=\"flow_m3d\").plot()\n",
     "ax.legend(bbox_to_anchor=(1.3, 1), title=\"Edge\")"
    ]
@@ -1405,7 +1405,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Model with allocation"
+    "# Model with allocation (user demand)"
    ]
   },
   {
@@ -1838,8 +1838,8 @@
     "df_allocation = pd.read_feather(datadir / \"allocation_example/results/allocation.arrow\")\n",
     "df_allocation_wide = df_allocation.pivot_table(\n",
     "    index=\"time\",\n",
-    "    columns=[\"user_node_id\", \"priority\"],\n",
-    "    values=[\"demand\", \"allocated\", \"abstracted\"],\n",
+    "    columns=[\"node_type\", \"node_id\", \"priority\"],\n",
+    "    values=[\"demand\", \"allocated\", \"realized\"],\n",
     ")\n",
     "df_allocation_wide = df_allocation_wide.loc[:, (df_allocation_wide != 0).any(axis=0)]\n",
     "\n",
@@ -1848,7 +1848,7 @@
     "\n",
     "df_allocation_wide[\"demand\"].plot(ax=axs[0], ls=\":\")\n",
     "df_allocation_wide[\"allocated\"].plot(ax=axs[1], ls=\"--\")\n",
-    "df_allocation_wide[\"abstracted\"].plot(ax=axs[2])\n",
+    "df_allocation_wide[\"realized\"].plot(ax=axs[2])\n",
     "\n",
     "fig.tight_layout()\n",
     "loc = plticker.MultipleLocator(2)\n",
@@ -1887,6 +1887,346 @@
     "ax.set_title(\"Basin levels\")\n",
     "ax.set_ylabel(\"level [m]\")"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Model with allocation (basin supply/demand)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Setup the nodes:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "xy = np.array(\n",
+    "    [\n",
+    "        (0.0, 0.0),  # 1: FlowBoundary\n",
+    "        (1.0, 0.0),  # 2: Basin\n",
+    "        (2.0, 0.0),  # 3: User\n",
+    "        (1.0, -1.0),  # 4: TargetLevel\n",
+    "        (2.0, -1.0),  # 5: Basin\n",
+    "    ]\n",
+    ")\n",
+    "node_xy = gpd.points_from_xy(x=xy[:, 0], y=xy[:, 1])\n",
+    "\n",
+    "node_type = [\"FlowBoundary\", \"Basin\", \"User\", \"TargetLevel\", \"Basin\"]\n",
+    "\n",
+    "# Make sure the feature id starts at 1: explicitly give an index.\n",
+    "node = ribasim.Node(\n",
+    "    df=gpd.GeoDataFrame(\n",
+    "        data={\n",
+    "            \"type\": node_type,\n",
+    "            \"subnetwork_id\": 5 * [2],\n",
+    "        },\n",
+    "        index=pd.Index(np.arange(len(xy)) + 1, name=\"fid\"),\n",
+    "        geometry=node_xy,\n",
+    "        crs=\"EPSG:28992\",\n",
+    "    )\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Setup the edges:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from_id = np.array([1, 2, 4, 3, 4])\n",
+    "to_id = np.array([2, 3, 2, 5, 5])\n",
+    "edge_type = [\"flow\", \"flow\", \"control\", \"flow\", \"control\"]\n",
+    "subnetwork_id = [2, None, None, None, None]\n",
+    "\n",
+    "lines = node.geometry_from_connectivity(from_id.tolist(), to_id.tolist())\n",
+    "edge = ribasim.Edge(\n",
+    "    df=gpd.GeoDataFrame(\n",
+    "        data={\n",
+    "            \"from_node_id\": from_id,\n",
+    "            \"to_node_id\": to_id,\n",
+    "            \"edge_type\": edge_type,\n",
+    "            \"subnetwork_id\": subnetwork_id,\n",
+    "        },\n",
+    "        geometry=lines,\n",
+    "        crs=\"EPSG:28992\",\n",
+    "    )\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Setup the basins:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "profile = pd.DataFrame(\n",
+    "    data={\"node_id\": [2, 2, 5, 5], \"area\": 1e3, \"level\": [0.0, 1.0, 0.0, 1.0]}\n",
+    ")\n",
+    "static = pd.DataFrame(\n",
+    "    data={\n",
+    "        \"node_id\": [5],\n",
+    "        \"drainage\": 0.0,\n",
+    "        \"potential_evaporation\": 0.0,\n",
+    "        \"infiltration\": 0.0,\n",
+    "        \"precipitation\": 0.0,\n",
+    "        \"urban_runoff\": 0.0,\n",
+    "    }\n",
+    ")\n",
+    "time = pd.DataFrame(\n",
+    "    data={\n",
+    "        \"node_id\": 2,\n",
+    "        \"time\": [\"2020-01-01 00:00:00\", \"2020-01-16 00:00:00\"],\n",
+    "        \"drainage\": 0.0,\n",
+    "        \"potential_evaporation\": 0.0,\n",
+    "        \"infiltration\": 0.0,\n",
+    "        \"precipitation\": [1e-6, 0.0],\n",
+    "        \"urban_runoff\": 0.0,\n",
+    "    },\n",
+    ")\n",
+    "\n",
+    "state = pd.DataFrame(data={\"node_id\": [2, 5], \"level\": 0.5})\n",
+    "basin = ribasim.Basin(profile=profile, static=static, time=time, state=state)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Setup the flow boundary:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "flow_boundary = ribasim.FlowBoundary(\n",
+    "    static=pd.DataFrame(data={\"node_id\": [1], \"flow_rate\": 1e-3})\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Setup allocation level control:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "target_level = ribasim.TargetLevel(\n",
+    "    static=pd.DataFrame(\n",
+    "        data={\"node_id\": [4], \"priority\": 1, \"min_level\": 1.0, \"max_level\": 1.5}\n",
+    "    )\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Setup the users:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "user = ribasim.User(\n",
+    "    static=pd.DataFrame(\n",
+    "        data={\n",
+    "            \"node_id\": [3],\n",
+    "            \"priority\": [2],\n",
+    "            \"demand\": [1.5e-3],\n",
+    "            \"return_factor\": [0.2],\n",
+    "            \"min_level\": [0.2],\n",
+    "        }\n",
+    "    )\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Setup the allocation:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "allocation = ribasim.Allocation(use_allocation=True, timestep=1e5)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Setup a model:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model = ribasim.Model(\n",
+    "    network=ribasim.Network(node=node, edge=edge),\n",
+    "    basin=basin,\n",
+    "    flow_boundary=flow_boundary,\n",
+    "    target_level=target_level,\n",
+    "    user=user,\n",
+    "    allocation=allocation,\n",
+    "    starttime=\"2020-01-01 00:00:00\",\n",
+    "    endtime=\"2020-02-01 00:00:00\",\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Let's take a look at the model:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.plot()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Write the model to a TOML and GeoPackage:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.write(datadir / \"target_level/ribasim.toml\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# | include: false\n",
+    "from subprocess import run\n",
+    "\n",
+    "run(\n",
+    "    [\n",
+    "        \"julia\",\n",
+    "        \"--project=../../core\",\n",
+    "        \"--eval\",\n",
+    "        f'using Ribasim; Ribasim.main(\"{datadir.as_posix()}/target_level/ribasim.toml\")',\n",
+    "    ],\n",
+    "    check=True,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Now run the model with `ribasim target_level/ribasim.toml`.\n",
+    "After running the model, read back the results:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "df_basin = pd.read_feather(datadir / \"target_level/results/basin.arrow\")\n",
+    "df_basin = df_basin[df_basin.node_id == 2]\n",
+    "df_basin_wide = df_basin.pivot_table(\n",
+    "    index=\"time\", columns=\"node_id\", values=[\"storage\", \"level\"]\n",
+    ")\n",
+    "ax = df_basin_wide[\"level\"].plot()\n",
+    "where_allocation = (\n",
+    "    df_basin_wide.index - df_basin_wide.index[0]\n",
+    ").total_seconds() % model.allocation.timestep == 0\n",
+    "where_allocation[0] = False\n",
+    "df_basin_wide[where_allocation][\"level\"].plot(\n",
+    "    style=\"o\",\n",
+    "    ax=ax,\n",
+    ")\n",
+    "ax.set_ylabel(\"level [m]\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "In the plot above, the line denotes the level of basin #2 over time and the dots denote the times at which allocation optimization was run, with intervals of $\\Delta t_{\\text{alloc}}$. The basin level is a piecewise linear function of time, with several stages explained below.\n",
+    "\n",
+    "Constants:\n",
+    "- $d$: user #3 demand,\n",
+    "- $\\phi$: precipitation rate on basin #2,\n",
+    "- $q$: flow of the level boundary."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Stages:\n",
+    "- In the first stage the user abstracts fully, so the net change of basin #2 is $q + \\phi - d$;\n",
+    "- In the second stage the basin takes precedence so the user doesn't abstract, hence the net change of basin #2 is $q + \\phi$;\n",
+    "-  In the third stage (and following stages) the basin has no longer a positive demand, since precipitation provides enough water to get the basin to its target level. The FlowBoundary flow gets fully allocated to the user, hence the net change of basin #2 is $\\phi$;\n",
+    "- In the fourth stage the basin enters its surplus stage, even though initially the level is below the maximum level. This is because the simulation anticipates that the current precipitation is going to bring the basin level over its maximum level. The net change of basin #2 is now $q + \\phi - d$;\n",
+    "- At the start of the fifth stage the precipitation stops, and so the user partly uses surplus water from the basin to fulfill its demand. The net change of basin #2 becomes $q - d$.\n",
+    "- In the final stage the basin is in a dynamical equilibrium, since the basin has no supply so the user abstracts precisely the flow from the level boundary."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": []
   }
  ],
  "metadata": {
diff --git a/python/ribasim_testmodels/ribasim_testmodels/allocation.py b/python/ribasim_testmodels/ribasim_testmodels/allocation.py
index ac379f2b3..b0282ca17 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/allocation.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/allocation.py
@@ -1674,7 +1674,7 @@ def target_level_model():
     from_id = np.array([1, 2, 4, 3, 4])
     to_id = np.array([2, 3, 2, 5, 5])
     edge_type = ["flow", "flow", "control", "flow", "control"]
-    allocation_network_id = [1, None, None, None, None]
+    subnetwork_id = [2, None, None, None, None]
 
     lines = node.geometry_from_connectivity(from_id.tolist(), to_id.tolist())
     edge = ribasim.Edge(
@@ -1683,7 +1683,7 @@ def target_level_model():
                 "from_node_id": from_id,
                 "to_node_id": to_id,
                 "edge_type": edge_type,
-                "allocation_network_id": allocation_network_id,
+                "subnetwork_id": subnetwork_id,
             },
             geometry=lines,
             crs="EPSG:28992",