From 6b7b3cca34312630aa65a5015ae409316e3e3878 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Mon, 8 Apr 2024 13:25:55 +0200
Subject: [PATCH 1/7] Move to 'quadrative_relative' objective function

---
 core/src/allocation_init.jl  | 104 -----------------------------------
 core/src/allocation_optim.jl |  92 ++++++-------------------------
 2 files changed, 16 insertions(+), 180 deletions(-)

diff --git a/core/src/allocation_init.jl b/core/src/allocation_init.jl
index d17ffc094..5f69d42b1 100644
--- a/core/src/allocation_init.jl
+++ b/core/src/allocation_init.jl
@@ -692,106 +692,6 @@ function add_constraints_conservation_basin!(
     return nothing
 end
 
-"""
-Minimizing |expr| can be achieved by introducing a new variable expr_abs
-and posing the following constraints:
-expr_abs >= expr
-expr_abs >= -expr
-"""
-function add_constraints_absolute_value!(
-    problem::JuMP.Model,
-    flow_per_node::Dict{NodeID, JuMP.VariableRef},
-    F_abs::JuMP.Containers.DenseAxisArray,
-    variable_type::String,
-)::Nothing
-    # Example demand
-    d = 2.0
-
-    node_ids = only(F_abs.axes)
-
-    # These constraints together make sure that F_abs_* acts as the absolute
-    # value F_abs_* = |x| where x = F-d (here for example d = 2)
-    base_name = "abs_positive_$variable_type"
-    problem[Symbol(base_name)] = JuMP.@constraint(
-        problem,
-        [node_id = node_ids],
-        F_abs[node_id] >= (flow_per_node[node_id] - d),
-        base_name = base_name
-    )
-    base_name = "abs_negative_$variable_type"
-    problem[Symbol(base_name)] = JuMP.@constraint(
-        problem,
-        [node_id = node_ids],
-        F_abs[node_id] >= -(flow_per_node[node_id] - d),
-        base_name = base_name
-    )
-
-    return nothing
-end
-
-"""
-Add constraints so that variables F_abs_user_demand act as the
-absolute value of the expression comparing flow to a UserDemand to its demand.
-"""
-function add_constraints_absolute_value_user_demand!(
-    problem::JuMP.Model,
-    p::Parameters,
-)::Nothing
-    (; graph) = p
-
-    F = problem[:F]
-    F_abs_user_demand = problem[:F_abs_user_demand]
-
-    flow_per_node = Dict(
-        node_id => F[(only(inflow_ids_allocation(graph, node_id)), node_id)] for
-        node_id in only(F_abs_user_demand.axes)
-    )
-
-    add_constraints_absolute_value!(
-        problem,
-        flow_per_node,
-        F_abs_user_demand,
-        "user_demand",
-    )
-
-    return nothing
-end
-
-"""
-Add constraints so that variables F_abs_level_demand act as the
-absolute value of the expression comparing flow to a basin to its demand.
-"""
-function add_constraints_absolute_value_level_demand!(problem::JuMP.Model)::Nothing
-    F_basin_in = problem[:F_basin_in]
-    F_abs_level_demand = problem[:F_abs_level_demand]
-    flow_per_node =
-        Dict(node_id => F_basin_in[node_id] for node_id in only(F_abs_level_demand.axes))
-
-    add_constraints_absolute_value!(problem, flow_per_node, F_abs_level_demand, "basin")
-
-    return nothing
-end
-
-"""
-Add constraints so that variables F_abs_flow_demand act as the
-absolute value of the expression comparing flow to a flow buffer to the flow demand.
-"""
-function add_constraints_absolute_value_flow_demand!(problem::JuMP.Model)::Nothing
-    F_flow_buffer_in = problem[:F_flow_buffer_in]
-    F_abs_flow_demand = problem[:F_abs_flow_demand]
-    flow_per_node = Dict(
-        node_id => F_flow_buffer_in[node_id] for node_id in only(F_abs_flow_demand.axes)
-    )
-
-    add_constraints_absolute_value!(
-        problem,
-        flow_per_node,
-        F_abs_flow_demand,
-        "flow_demand",
-    )
-    return nothing
-end
-
 """
 Add the fractional flow constraints to the allocation problem.
 The constraint indices are allocation edges over a fractional flow node.
@@ -929,10 +829,6 @@ function allocation_problem(
     add_constraints_conservation_flow_demand!(problem, p, allocation_network_id)
     add_constraints_conservation_subnetwork!(problem, p, allocation_network_id)
 
-    add_constraints_absolute_value_user_demand!(problem, p)
-    add_constraints_absolute_value_flow_demand!(problem)
-    add_constraints_absolute_value_level_demand!(problem)
-
     add_constraints_capacity!(problem, capacity, p, allocation_network_id)
     add_constraints_source!(problem, p, allocation_network_id)
     add_constraints_user_source!(problem, p, allocation_network_id)
diff --git a/core/src/allocation_optim.jl b/core/src/allocation_optim.jl
index b42393f75..9920799a1 100644
--- a/core/src/allocation_optim.jl
+++ b/core/src/allocation_optim.jl
@@ -1,67 +1,16 @@
 @enumx OptimizationType internal_sources collect_demands allocate
 
-"""
-Add a term to the objective function given by the objective type,
-depending in the provided flow variable and the associated demand.
-"""
 function add_objective_term!(
+    ex::JuMP.QuadExpr,
     demand::Float64,
-    constraint_abs_positive::Union{JuMP.ConstraintRef, Nothing} = nothing,
-    constraint_abs_negative::Union{JuMP.ConstraintRef, Nothing} = nothing,
-)::Nothing
-    # Objective function ∑ |F - d|
-    JuMP.set_normalized_rhs(constraint_abs_positive, -demand)
-    JuMP.set_normalized_rhs(constraint_abs_negative, demand)
-    return nothing
-end
-
-"""
-Add a term to the expression of the objective function corresponding to
-the demand of a UserDemand.
-"""
-function add_user_demand_term!(
-    edge::Tuple{NodeID, NodeID},
-    demand::Float64,
-    problem::JuMP.Model,
+    F::JuMP.VariableRef,
 )::Nothing
-    node_id_user_demand = edge[2]
-
-    constraint_abs_positive = problem[:abs_positive_user_demand][node_id_user_demand]
-    constraint_abs_negative = problem[:abs_negative_user_demand][node_id_user_demand]
-
-    add_objective_term!(demand, constraint_abs_positive, constraint_abs_negative)
-end
-
-"""
-Add a term to the expression of the objective function corresponding to
-the demand of a node with a a flow demand.
-"""
-function add_flow_demand_term!(
-    edge::Tuple{NodeID, NodeID},
-    demand::Float64,
-    problem::JuMP.Model,
-)::Nothing
-    node_id_flow_demand = edge[2]
-
-    constraint_abs_positive = problem[:abs_positive_flow_demand][node_id_flow_demand]
-    constraint_abs_negative = problem[:abs_negative_flow_demand][node_id_flow_demand]
-
-    add_objective_term!(demand, constraint_abs_positive, constraint_abs_negative)
-end
-
-"""
-Add a term to the expression of the objective function corresponding to
-the demand of a basin.
-"""
-function add_basin_term!(problem::JuMP.Model, demand::Float64, node_id::NodeID)::Nothing
-    constraint_abs_positive = get(problem[:abs_positive_basin], node_id)
-    constraint_abs_negative = get(problem[:abs_negative_basin], node_id)
-
-    if isnothing(constraint_abs_positive)
+    if demand ≈ 0
         return
     end
-
-    add_objective_term!(demand, constraint_abs_positive, constraint_abs_negative)
+    JuMP.add_to_expression!(ex, 1.0 / demand^2, F, F)
+    JuMP.add_to_expression!(ex, -2.0 / demand, F)
+    JuMP.add_to_expression!(ex, 1.0)
     return nothing
 end
 
@@ -81,29 +30,17 @@ function set_objective_priority!(
     (; node_id, demand_reduced) = user_demand
     (; main_network_connections, subnetwork_demands) = allocation
     edge_ids = graph[].edge_ids[allocation_network_id]
+    F = problem[:F]
 
-    ex = JuMP.AffExpr()
-
-    F_abs_user_demand = problem[:F_abs_user_demand]
-    F_abs_level_demand = problem[:F_abs_level_demand]
-    F_abs_flow_demand = problem[:F_abs_flow_demand]
-
-    if !isempty(only(F_abs_user_demand.axes))
-        ex += sum(F_abs_user_demand)
-    end
-    if !isempty(only(F_abs_level_demand.axes))
-        ex += sum(F_abs_level_demand)
-    end
-    if !isempty(only(F_abs_flow_demand.axes))
-        ex += sum(F_abs_flow_demand)
-    end
+    ex = JuMP.QuadExpr()
 
     # Terms for subnetworks as UserDemand
     if is_main_network(allocation_network_id)
         for connections_subnetwork in main_network_connections
             for connection in connections_subnetwork
                 d = subnetwork_demands[connection][priority_idx]
-                add_user_demand_term!(connection, d, problem)
+                F_inlet = F[connection]
+                add_objective_term!(ex, d, F_inlet)
             end
         end
     end
@@ -116,7 +53,8 @@ function set_objective_priority!(
             # UserDemand
             user_demand_idx = findsorted(node_id, to_node_id)
             d = demand_reduced[user_demand_idx, priority_idx]
-            add_user_demand_term!(edge_id, d, problem)
+            F_ud = F[edge_id]
+            add_objective_term!(ex, d, F_ud)
         else
             has_demand, demand_node_id =
                 has_external_demand(graph, to_node_id, :flow_demand)
@@ -127,8 +65,9 @@ function set_objective_priority!(
                     priority_idx == flow_priority_idx ?
                     flow_demand.demand[findsorted(flow_demand.node_id, demand_node_id)] :
                     0.0
+                F_fd = F[edge_id]
 
-                add_flow_demand_term!(edge_id, d, problem)
+                add_objective_term!(ex, d, F_fd)
             end
         end
     end
@@ -142,7 +81,8 @@ function set_objective_priority!(
             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!(problem, d, node_id)
+        F_ld = F_basin_in[node_id]
+        add_objective_term!(ex, d, F_ld)
     end
 
     new_objective = JuMP.@expression(problem, ex)

From ddde12f7a615ce87cbdcd1035255b8a34ffa29ec Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Mon, 8 Apr 2024 17:00:45 +0200
Subject: [PATCH 2/7] Fix most tests

---
 core/src/allocation_optim.jl |  2 +-
 core/test/allocation_test.jl | 96 +++++++++++++++++++++++-------------
 2 files changed, 62 insertions(+), 36 deletions(-)

diff --git a/core/src/allocation_optim.jl b/core/src/allocation_optim.jl
index 9920799a1..0f2378646 100644
--- a/core/src/allocation_optim.jl
+++ b/core/src/allocation_optim.jl
@@ -5,7 +5,7 @@ function add_objective_term!(
     demand::Float64,
     F::JuMP.VariableRef,
 )::Nothing
-    if demand ≈ 0
+    if abs(demand) < 1e-5
         return
     end
     JuMP.add_to_expression!(ex, 1.0 / demand^2, F, F)
diff --git a/core/test/allocation_test.jl b/core/test/allocation_test.jl
index d804d8020..9a5603c62 100644
--- a/core/test/allocation_test.jl
+++ b/core/test/allocation_test.jl
@@ -57,15 +57,23 @@ end
     config = Ribasim.Config(toml_path)
     model = Ribasim.run(config)
     @test successful_retcode(model)
-    problem = model.integrator.p.allocation.allocation_models[1].problem
+    (; p) = model.integrator
+    (; user_demand) = p
+    problem = p.allocation.allocation_models[1].problem
     objective = JuMP.objective_function(problem)
-    @test objective isa JuMP.AffExpr # Affine expression
-    @test :F_abs_user_demand in keys(problem.obj_dict)
+    @test objective isa JuMP.QuadExpr # Quadratic expression
     F = problem[:F]
-    F_abs_user_demand = problem[:F_abs_user_demand]
 
-    @test objective.terms[F_abs_user_demand[NodeID(:UserDemand, 5)]] == 1.0
-    @test objective.terms[F_abs_user_demand[NodeID(:UserDemand, 6)]] == 1.0
+    to_user_5 = F[(NodeID(:Basin, 4), NodeID(:UserDemand, 5))]
+    to_user_6 = F[(NodeID(:Basin, 4), NodeID(:UserDemand, 6))]
+
+    @test objective.aff.constant == 2.0
+    @test objective.aff.terms[to_user_5] ≈ -2.0 / user_demand.demand[1]
+    @test objective.aff.terms[to_user_6] ≈ -2.0 / user_demand.demand[2]
+    @test objective.terms[JuMP.UnorderedPair(to_user_5, to_user_5)] ≈
+          1 / user_demand.demand[1]^2
+    @test objective.terms[JuMP.UnorderedPair(to_user_6, to_user_6)] ≈
+          1 / user_demand.demand[2]^2
 end
 
 @testitem "Allocation with controlled fractional flow" begin
@@ -173,14 +181,6 @@ end
         (NodeID(:Basin, 10), NodeID(:Pump, 38)),
     ] ⊆ allocation_edges_main_network
 
-    # Subnetworks interpreted as user_demands require variables and constraints to
-    # support absolute value expressions in the objective function
-    allocation_model_main_network = Ribasim.get_allocation_model(p, Int32(1))
-    problem = allocation_model_main_network.problem
-    @test problem[:F_abs_user_demand].axes[1] == NodeID.(:Pump, [11, 24, 38])
-    @test problem[:abs_positive_user_demand].axes[1] == NodeID.(:Pump, [11, 24, 38])
-    @test problem[:abs_negative_user_demand].axes[1] == NodeID.(:Pump, [11, 24, 38])
-
     # In each subnetwork, the connection from the main network to the subnetwork is
     # interpreted as a source
     @test Ribasim.get_allocation_model(p, Int32(3)).problem[:source].axes[1] ==
@@ -223,33 +223,53 @@ end
 
     # See the difference between these values here and in
     # "subnetworks_with_sources"
-    @test subnetwork_demands[(NodeID(:Basin, 2), NodeID(:Pump, 11))] ≈ [4.0, 4.0, 0.0]
+    @test all(
+        isapprox.(
+            subnetwork_demands[(NodeID(:Basin, 2), NodeID(:Pump, 11))],
+            [4.0, 4.0, 0.0],
+            atol = 1e-4,
+        ),
+    )
     @test subnetwork_demands[(NodeID(:Basin, 6), NodeID(:Pump, 24))] ≈ [0.004, 0.0, 0.0]
-    @test subnetwork_demands[(NodeID(:Basin, 10), NodeID(:Pump, 38))][1:2] ≈ [0.001, 0.002]
+    @test subnetwork_demands[(NodeID(:Basin, 10), NodeID(:Pump, 38))][1:2] ≈ [0.004, 0.001]
 
     # Solving for the main network, containing subnetworks as UserDemands
     allocation_model = allocation_models[1]
     (; problem) = allocation_model
-    Ribasim.allocate!(p, allocation_model, t, u, OptimizationType.allocate)
+    Ribasim.allocate_priority!(allocation_model, u, p, t, 1, OptimizationType.allocate)
 
     # Main network objective function
+    F = problem[:F]
     objective = JuMP.objective_function(problem)
-    objective_variables = keys(objective.terms)
-    F_abs_user_demand = problem[:F_abs_user_demand]
-    @test F_abs_user_demand[NodeID(:Pump, 11)] ∈ objective_variables
-    @test F_abs_user_demand[NodeID(:Pump, 24)] ∈ objective_variables
-    @test F_abs_user_demand[NodeID(:Pump, 38)] ∈ objective_variables
+    objective_edges = keys(objective.terms)
+    F_1 = F[(NodeID(:Basin, 2), NodeID(:Pump, 11))]
+    F_2 = F[(NodeID(:Basin, 6), NodeID(:Pump, 24))]
+    F_3 = F[(NodeID(:Basin, 10), NodeID(:Pump, 38))]
+    @test JuMP.UnorderedPair(F_1, F_1) ∈ objective_variables
+    @test JuMP.UnorderedPair(F_2, F_2) ∈ objective_variables
+    @test JuMP.UnorderedPair(F_3, F_3) ∈ objective_variables
 
     # Running full allocation algorithm
     Ribasim.set_flow!(graph, NodeID(:FlowBoundary, 1), NodeID(:Basin, 2), 4.5)
     u = ComponentVector(; storage = zeros(length(p.basin.node_id)))
     Ribasim.update_allocation!((; p, t, u))
 
-    @test subnetwork_allocateds[NodeID(:Basin, 2), NodeID(:Pump, 11)] ≈
-          [4.0, 0.49500000, 0.0]
+    @test all(
+        isapprox.(
+            subnetwork_allocateds[NodeID(:Basin, 2), NodeID(:Pump, 11)],
+            [4.0, 0.49100000, 0.0],
+            atol = 1e-4,
+        ),
+    )
     @test subnetwork_allocateds[NodeID(:Basin, 6), NodeID(:Pump, 24)] ≈
           [0.00399999999, 0.0, 0.0]
-    @test subnetwork_allocateds[NodeID(:Basin, 10), NodeID(:Pump, 38)] ≈ [0.001, 0.0, 0.0]
+    @test all(
+        isapprox.(
+            subnetwork_allocateds[NodeID(:Basin, 10), NodeID(:Pump, 38)],
+            [0.004, 0.001, 0.0],
+            atol = 1e-4,
+        ),
+    )
 
     # Test for existence of edges in allocation flow record
     allocation_flow = DataFrame(record_flow)
@@ -261,7 +281,7 @@ end
     )
     @test all(allocation_flow.edge_exists)
 
-    @test user_demand.allocated[2, :] ≈ [4.0, 0.0, 0.0]
+    @test user_demand.allocated[2, :] ≈ [4.0, 0.0, 3.6]
     @test user_demand.allocated[7, :] ≈ [0.001, 0.0, 0.0]
 end
 
@@ -292,7 +312,8 @@ end
 
     # Collecting demands
     u = ComponentVector(; storage = zeros(length(basin.node_id)))
-    for allocation_model in allocation_models[2:end]
+    for (i, allocation_model) in enumerate(allocation_models[2:end])
+        @show i
         Ribasim.allocate!(p, allocation_model, t, u, OptimizationType.internal_sources)
         Ribasim.allocate!(p, allocation_model, t, u, OptimizationType.collect_demands)
     end
@@ -300,7 +321,13 @@ end
     # See the difference between these values here and in
     # "allocation with main network optimization problem", internal sources
     # lower the subnetwork demands
-    @test subnetwork_demands[(NodeID(:Basin, 2), NodeID(:Pump, 11))] ≈ [3.1, 4.0, 0.0]
+    @test all(
+        isapprox.(
+            subnetwork_demands[(NodeID(:Basin, 2), NodeID(:Pump, 11))],
+            [2.29, 4.0, 0.0],
+            atol = 1e-4,
+        ),
+    )
     @test subnetwork_demands[(NodeID(:Basin, 6), NodeID(:Pump, 24))] ≈ [0.004, 0.0, 0.0]
     @test subnetwork_demands[(NodeID(:Basin, 10), NodeID(:Pump, 38))][1:2] ≈ [0.001, 0.001]
 end
@@ -435,9 +462,8 @@ end
         optimization_type,
     )
     objective = JuMP.objective_function(problem)
-    @test F_abs_flow_demand[node_id_with_flow_demand] in keys(objective.terms)
     # Reduced demand
-    @test flow_demand.demand[1] == flow_demand.demand_itp[1](t) - 0.001
+    @test flow_demand.demand[1] ≈ flow_demand.demand_itp[1](t) - 0.001
     @test JuMP.normalized_rhs(constraint_flow_out) == Inf
 
     ## Priority 2
@@ -450,9 +476,9 @@ end
         optimization_type,
     )
     # No demand left
-    @test flow_demand.demand[1] ≈ 0.0
+    @test flow_demand.demand[1] < 1e-10
     # Allocated
-    @test JuMP.value(only(F_flow_buffer_in)) == 0.001
+    @test JuMP.value(only(F_flow_buffer_in)) ≈ 0.001
     @test JuMP.normalized_rhs(constraint_flow_out) == 0.0
 
     ## Priority 3
@@ -468,9 +494,9 @@ end
     # The flow from the source is used up in previous priorities
     @test JuMP.value(F[(NodeID(NodeType.LevelBoundary, 1), node_id_with_flow_demand)]) == 0
     # So flow from the flow buffer is used for UserDemand #4
-    @test JuMP.value(F_flow_buffer_out[node_id_with_flow_demand]) == 0.001
+    @test JuMP.value(F_flow_buffer_out[node_id_with_flow_demand]) ≈ 0.001
     # Flow taken from buffer
-    @test JuMP.value(only(F_flow_buffer_out)) == user_demand.demand_itp[1][3](t)
+    @test JuMP.value(only(F_flow_buffer_out)) ≈ user_demand.demand_itp[1][3](t)
     # No flow coming from level boundary
     @test JuMP.value(F[(only(level_boundary.node_id), node_id_with_flow_demand)]) == 0
 
@@ -486,7 +512,7 @@ end
     # Get demand from buffers
     d = user_demand.demand_itp[3][4](t)
     @assert JuMP.value(F[(NodeID(NodeType.UserDemand, 4), NodeID(NodeType.Basin, 7))]) +
-            JuMP.value(F[(NodeID(NodeType.UserDemand, 6), NodeID(NodeType.Basin, 7))]) == d
+            JuMP.value(F[(NodeID(NodeType.UserDemand, 6), NodeID(NodeType.Basin, 7))]) ≈ d
 end
 
 @testitem "flow_demand_with_max_flow_rate" begin

From cb4a618f53da5f479207e9325db68463820df2c3 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Tue, 9 Apr 2024 11:52:48 +0200
Subject: [PATCH 3/7] Fix existing tests

---
 core/src/allocation_optim.jl | 18 +++++++++---
 core/test/allocation_test.jl | 57 ++++++++++++++++++++++++------------
 2 files changed, 53 insertions(+), 22 deletions(-)

diff --git a/core/src/allocation_optim.jl b/core/src/allocation_optim.jl
index 0f2378646..9c112127d 100644
--- a/core/src/allocation_optim.jl
+++ b/core/src/allocation_optim.jl
@@ -1,16 +1,26 @@
 @enumx OptimizationType internal_sources collect_demands allocate
 
+"""
+Add an objective term (1 - flow/demand)^2. If the absolute
+value of the demand is very small, this would lead to huge coefficients,
+so in that case a term of the form (flow - demand)^2 is used.
+"""
 function add_objective_term!(
     ex::JuMP.QuadExpr,
     demand::Float64,
     F::JuMP.VariableRef,
 )::Nothing
     if abs(demand) < 1e-5
-        return
+        # Error term (F - d)^2
+        JuMP.add_to_expression!(ex, 1.0, F, F)
+        JuMP.add_to_expression!(ex, -2.0 * demand, F)
+        JuMP.add_to_expression!(ex, demand^2)
+    else
+        # Error term (1 - F/d)^2
+        JuMP.add_to_expression!(ex, 1.0 / demand^2, F, F)
+        JuMP.add_to_expression!(ex, -2.0 / demand, F)
+        JuMP.add_to_expression!(ex, 1.0)
     end
-    JuMP.add_to_expression!(ex, 1.0 / demand^2, F, F)
-    JuMP.add_to_expression!(ex, -2.0 / demand, F)
-    JuMP.add_to_expression!(ex, 1.0)
     return nothing
 end
 
diff --git a/core/test/allocation_test.jl b/core/test/allocation_test.jl
index 9a5603c62..e720c042d 100644
--- a/core/test/allocation_test.jl
+++ b/core/test/allocation_test.jl
@@ -231,7 +231,13 @@ end
         ),
     )
     @test subnetwork_demands[(NodeID(:Basin, 6), NodeID(:Pump, 24))] ≈ [0.004, 0.0, 0.0]
-    @test subnetwork_demands[(NodeID(:Basin, 10), NodeID(:Pump, 38))][1:2] ≈ [0.004, 0.001]
+    @test all(
+        isapprox.(
+            subnetwork_demands[(NodeID(:Basin, 10), NodeID(:Pump, 38))][1:2],
+            [0.001, 0.002],
+            rtol = 1e-4,
+        ),
+    )
 
     # Solving for the main network, containing subnetworks as UserDemands
     allocation_model = allocation_models[1]
@@ -245,9 +251,9 @@ end
     F_1 = F[(NodeID(:Basin, 2), NodeID(:Pump, 11))]
     F_2 = F[(NodeID(:Basin, 6), NodeID(:Pump, 24))]
     F_3 = F[(NodeID(:Basin, 10), NodeID(:Pump, 38))]
-    @test JuMP.UnorderedPair(F_1, F_1) ∈ objective_variables
-    @test JuMP.UnorderedPair(F_2, F_2) ∈ objective_variables
-    @test JuMP.UnorderedPair(F_3, F_3) ∈ objective_variables
+    @test JuMP.UnorderedPair(F_1, F_1) ∈ objective_edges
+    @test JuMP.UnorderedPair(F_2, F_2) ∈ objective_edges
+    @test JuMP.UnorderedPair(F_3, F_3) ∈ objective_edges
 
     # Running full allocation algorithm
     Ribasim.set_flow!(graph, NodeID(:FlowBoundary, 1), NodeID(:Basin, 2), 4.5)
@@ -257,17 +263,20 @@ end
     @test all(
         isapprox.(
             subnetwork_allocateds[NodeID(:Basin, 2), NodeID(:Pump, 11)],
-            [4.0, 0.49100000, 0.0],
+            [4.0, 0.493, 0.0],
             atol = 1e-4,
         ),
     )
-    @test subnetwork_allocateds[NodeID(:Basin, 6), NodeID(:Pump, 24)] ≈
-          [0.00399999999, 0.0, 0.0]
+    @test isapprox(
+        subnetwork_allocateds[NodeID(:Basin, 6), NodeID(:Pump, 24)],
+        [0.004, 0.0, 0.0],
+        rtol = 1e-3,
+    )
     @test all(
         isapprox.(
             subnetwork_allocateds[NodeID(:Basin, 10), NodeID(:Pump, 38)],
-            [0.004, 0.001, 0.0],
-            atol = 1e-4,
+            [0.001, 0.002, 0.0],
+            rtol = 1e-3,
         ),
     )
 
@@ -281,8 +290,8 @@ end
     )
     @test all(allocation_flow.edge_exists)
 
-    @test user_demand.allocated[2, :] ≈ [4.0, 0.0, 3.6]
-    @test user_demand.allocated[7, :] ≈ [0.001, 0.0, 0.0]
+    @test user_demand.allocated[2, :] ≈ [4.0, 0.0, 0.0]
+    @test all(isapprox.(user_demand.allocated[7, :], [0.001, 0.0, 0.0], atol = 1e-5))
 end
 
 @testitem "subnetworks with sources" begin
@@ -324,7 +333,7 @@ end
     @test all(
         isapprox.(
             subnetwork_demands[(NodeID(:Basin, 2), NodeID(:Pump, 11))],
-            [2.29, 4.0, 0.0],
+            [3.1, 4.0, 0.0],
             atol = 1e-4,
         ),
     )
@@ -463,7 +472,7 @@ end
     )
     objective = JuMP.objective_function(problem)
     # Reduced demand
-    @test flow_demand.demand[1] ≈ flow_demand.demand_itp[1](t) - 0.001
+    @test isapprox(flow_demand.demand[1], flow_demand.demand_itp[1](t) - 0.001, rtol = 1e-4)
     @test JuMP.normalized_rhs(constraint_flow_out) == Inf
 
     ## Priority 2
@@ -478,7 +487,7 @@ end
     # No demand left
     @test flow_demand.demand[1] < 1e-10
     # Allocated
-    @test JuMP.value(only(F_flow_buffer_in)) ≈ 0.001
+    @test isapprox(JuMP.value(only(F_flow_buffer_in)), 0.001, rtol = 1e-4)
     @test JuMP.normalized_rhs(constraint_flow_out) == 0.0
 
     ## Priority 3
@@ -494,9 +503,17 @@ end
     # The flow from the source is used up in previous priorities
     @test JuMP.value(F[(NodeID(NodeType.LevelBoundary, 1), node_id_with_flow_demand)]) == 0
     # So flow from the flow buffer is used for UserDemand #4
-    @test JuMP.value(F_flow_buffer_out[node_id_with_flow_demand]) ≈ 0.001
+    @test isapprox(
+        JuMP.value(F_flow_buffer_out[node_id_with_flow_demand]),
+        0.001,
+        rtol = 1e-4,
+    )
     # Flow taken from buffer
-    @test JuMP.value(only(F_flow_buffer_out)) ≈ user_demand.demand_itp[1][3](t)
+    @test isapprox(
+        JuMP.value(only(F_flow_buffer_out)),
+        user_demand.demand_itp[1][3](t),
+        rtol = 1e-4,
+    )
     # No flow coming from level boundary
     @test JuMP.value(F[(only(level_boundary.node_id), node_id_with_flow_demand)]) == 0
 
@@ -511,8 +528,12 @@ end
     )
     # Get demand from buffers
     d = user_demand.demand_itp[3][4](t)
-    @assert JuMP.value(F[(NodeID(NodeType.UserDemand, 4), NodeID(NodeType.Basin, 7))]) +
-            JuMP.value(F[(NodeID(NodeType.UserDemand, 6), NodeID(NodeType.Basin, 7))]) ≈ d
+    @assert isapprox(
+        JuMP.value(F[(NodeID(NodeType.UserDemand, 4), NodeID(NodeType.Basin, 7))]) +
+        JuMP.value(F[(NodeID(NodeType.UserDemand, 6), NodeID(NodeType.Basin, 7))]),
+        d,
+        rtol = 1e-4,
+    )
 end
 
 @testitem "flow_demand_with_max_flow_rate" begin

From 1cd9c8efbe2302d1e6572ac10eb2ef7f4a816584 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Tue, 9 Apr 2024 13:23:56 +0200
Subject: [PATCH 4/7] Start of test model

---
 .../ribasim_testmodels/__init__.py            |  2 +
 .../ribasim_testmodels/allocation.py          | 90 +++++++++++++++++++
 2 files changed, 92 insertions(+)

diff --git a/python/ribasim_testmodels/ribasim_testmodels/__init__.py b/python/ribasim_testmodels/ribasim_testmodels/__init__.py
index e4dd9dbad..888e5778b 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/__init__.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/__init__.py
@@ -7,6 +7,7 @@
 import ribasim_testmodels
 from ribasim_testmodels.allocation import (
     allocation_example_model,
+    fair_distribution_model,
     flow_demand_model,
     fractional_flow_subnetwork_model,
     level_demand_model,
@@ -65,6 +66,7 @@
     "bucket_model",
     "discrete_control_of_pid_control_model",
     "dutch_waterways_model",
+    "fair_distribution_model",
     "flow_boundary_time_model",
     "flow_condition_model",
     "flow_demand_model",
diff --git a/python/ribasim_testmodels/ribasim_testmodels/allocation.py b/python/ribasim_testmodels/ribasim_testmodels/allocation.py
index 1e102e8c0..df2f8383e 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/allocation.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/allocation.py
@@ -1029,3 +1029,93 @@ def linear_resistance_demand_model():
     model.edge.add(model.flow_demand[4], model.linear_resistance[2])
 
     return model
+
+
+def fair_distribution_model():
+    model = Model(
+        starttime="2020-01-01 00:00:00",
+        endtime="2021-01-01 00:00:00",
+        crs="EPSG:28992",
+        allocation=Allocation(use_allocation=True),
+    )
+
+    model.flow_boundary.add(
+        Node(1, Point(0, 0), subnetwork_id=1),
+        [
+            flow_boundary.Time(
+                time=pd.date_range(start="2020-01", end="2020-05", freq="MS"),
+                flow_rate=np.arange(0, 12.5, 2.5),
+            )
+        ],
+    )
+
+    model.linear_resistance.add(
+        Node(
+            2,
+            Point(1, 0),
+            subnetwork_id=1,
+        ),
+        [linear_resistance.Static(flow_rate=[9.0], max_flow_rate=9.0)],
+    )
+
+    model.basin.add(
+        Node(3, Point(2, 0), subnetwork_id=1),
+        [basin.Profile(area=1e3, level=[0.0, 1.0]), basin.State(level=[1.0])],
+    )
+
+    model.linear_resistance.add(
+        Node(4, Point(3, 0), subnetwork_id=1),
+        [linear_resistance.Static(resistance=[1.0])],
+    )
+
+    model.basin.add(
+        Node(5, Point(4, 0), subnetwork_id=1),
+        [basin.Profile(area=1e3, level=[0.0, 1.0]), basin.State(level=[1.0])],
+    )
+
+    model.user_demand.add(
+        Node(6, Point(2, 1), subnetwork_id=1),
+        [
+            user_demand.Static(
+                priority=[1], demand=1.0, return_factor=1.0, min_level=0.2
+            )
+        ],
+    )
+
+    model.user_demand.add(
+        Node(7, Point(2, -1), subnetwork_id=1),
+        [
+            user_demand.Static(
+                priority=[1], demand=1.0, return_factor=1.0, min_level=0.2
+            )
+        ],
+    )
+
+    model.user_demand.add(
+        Node(8, Point(4, 1), subnetwork_id=1),
+        [
+            user_demand.Static(
+                priority=[1], demand=1.0, return_factor=1.0, min_level=0.2
+            )
+        ],
+    )
+
+    model.user_demand.add(
+        Node(9, Point(4, -1), subnetwork_id=1),
+        [
+            user_demand.Static(
+                priority=[1], demand=1.0, return_factor=1.0, min_level=0.2
+            )
+        ],
+    )
+
+    model.edge.add(model.flow_boundary[1], model.pump[2])
+    model.edge.add(model.pump[2], model.basin[3])
+    model.edge.add(model.basin[3], model.linear_resistance[4])
+    model.edge.add(model.linear_resistance[4], model.basin[5])
+    model.edge.add(model.basin[3], model.user_demand[6])
+    model.edge.add(model.basin[3], model.user_demand[7])
+    model.edge.add(model.basin[5], model.user_demand[8])
+    model.edge.add(model.basin[5], model.user_demand[9])
+
+    return model

From 5996b965981889a59151b7de1494f6c590cdad74 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 10 Apr 2024 08:34:22 +0200
Subject: [PATCH 5/7] Fix main network source bug + fix tests

---
 core/src/allocation_init.jl                   | 54 +-----------------
 core/src/allocation_optim.jl                  | 13 +++--
 core/test/allocation_test.jl                  | 57 +++++++++++++------
 .../ribasim_testmodels/allocation.py          | 29 ++++++----
 4 files changed, 68 insertions(+), 85 deletions(-)

diff --git a/core/src/allocation_init.jl b/core/src/allocation_init.jl
index 5f69d42b1..0f16cb2b4 100644
--- a/core/src/allocation_init.jl
+++ b/core/src/allocation_init.jl
@@ -5,7 +5,8 @@ function find_subnetwork_connections!(p::Parameters)::Nothing
     (; subnetwork_demands, subnetwork_allocateds) = allocation
     for node_id in graph[].node_ids[1]
         for outflow_id in outflow_ids(graph, node_id)
-            if graph[outflow_id].allocation_network_id != 1
+            if (graph[outflow_id].allocation_network_id != 1) |
+               (graph[node_id, outflow_id].allocation_network_id_source == 1)
                 main_network_source_edges =
                     get_main_network_connections(p, graph[outflow_id].allocation_network_id)
                 edge = (node_id, outflow_id)
@@ -405,55 +406,6 @@ function add_variables_flow_buffer!(
     return nothing
 end
 
-"""
-Certain allocation distribution types use absolute values in the objective function.
-Since most optimization packages do not support the absolute value function directly,
-New variables are introduced that act as the absolute value of an expression by
-posing the appropriate constraints.
-"""
-function add_variables_absolute_value!(
-    problem::JuMP.Model,
-    p::Parameters,
-    allocation_network_id::Int32,
-)::Nothing
-    (; graph, allocation) = p
-    (; main_network_connections) = allocation
-
-    node_ids = graph[].node_ids[allocation_network_id]
-    node_ids_user_demand = NodeID[]
-    node_ids_level_demand = NodeID[]
-    node_ids_flow_demand = NodeID[]
-
-    for node_id in node_ids
-        type = node_id.type
-        if type == NodeType.UserDemand
-            push!(node_ids_user_demand, node_id)
-        elseif type == NodeType.Basin
-            push!(node_ids_level_demand, node_id)
-        elseif has_external_demand(graph, node_id, :flow_demand)[1]
-            push!(node_ids_flow_demand, node_id)
-        end
-    end
-
-    # For the main network, connections to subnetworks are treated as UserDemands
-    if is_main_network(allocation_network_id)
-        for connections_subnetwork in main_network_connections
-            for connection in connections_subnetwork
-                push!(node_ids_user_demand, connection[2])
-            end
-        end
-    end
-
-    problem[:F_abs_user_demand] =
-        JuMP.@variable(problem, F_abs_user_demand[node_id = node_ids_user_demand])
-    problem[:F_abs_level_demand] =
-        JuMP.@variable(problem, F_abs_level_demand[node_id = node_ids_level_demand])
-    problem[:F_abs_flow_demand] =
-        JuMP.@variable(problem, F_abs_flow_demand[node_id = node_ids_flow_demand])
-
-    return nothing
-end
-
 """
 Add the flow capacity constraints to the allocation problem.
 Only finite capacities get a constraint.
@@ -468,7 +420,6 @@ function add_constraints_capacity!(
     p::Parameters,
     allocation_network_id::Int32,
 )::Nothing
-    (; graph) = p
     main_network_source_edges = get_main_network_connections(p, allocation_network_id)
     F = problem[:F]
     edge_ids_finite_capacity = Tuple{NodeID, NodeID}[]
@@ -821,7 +772,6 @@ function allocation_problem(
     # Add variables to problem
     add_variables_flow!(problem, p, allocation_network_id)
     add_variables_basin!(problem, p, allocation_network_id)
-    add_variables_absolute_value!(problem, p, allocation_network_id)
     add_variables_flow_buffer!(problem, p, allocation_network_id)
 
     # Add constraints to problem
diff --git a/core/src/allocation_optim.jl b/core/src/allocation_optim.jl
index 9c112127d..07bc1ff29 100644
--- a/core/src/allocation_optim.jl
+++ b/core/src/allocation_optim.jl
@@ -16,10 +16,10 @@ function add_objective_term!(
         JuMP.add_to_expression!(ex, -2.0 * demand, F)
         JuMP.add_to_expression!(ex, demand^2)
     else
-        # Error term (1 - F/d)^2
-        JuMP.add_to_expression!(ex, 1.0 / demand^2, F, F)
-        JuMP.add_to_expression!(ex, -2.0 / demand, F)
-        JuMP.add_to_expression!(ex, 1.0)
+        # Error term d*(1 - F/d)^2
+        JuMP.add_to_expression!(ex, 1.0 / demand, F, F)
+        JuMP.add_to_expression!(ex, -2.0, F)
+        JuMP.add_to_expression!(ex, demand)
     end
     return nothing
 end
@@ -46,7 +46,7 @@ function set_objective_priority!(
 
     # Terms for subnetworks as UserDemand
     if is_main_network(allocation_network_id)
-        for connections_subnetwork in main_network_connections
+        for connections_subnetwork in main_network_connections[2:end]
             for connection in connections_subnetwork
                 d = subnetwork_demands[connection][priority_idx]
                 F_inlet = F[connection]
@@ -209,7 +209,8 @@ function set_initial_capacities_source!(
     for edge_id in edge_ids
         if graph[edge_id...].allocation_network_id_source == allocation_network_id
             # If it is a source edge for this allocation problem
-            if edge_id ∉ main_network_source_edges
+            if (edge_id ∉ main_network_source_edges) |
+               is_main_network(allocation_network_id)
                 # Reset the source to the current flow from the physical layer.
                 source_capacity = get_flow(graph, edge_id..., 0)
                 JuMP.set_normalized_rhs(
diff --git a/core/test/allocation_test.jl b/core/test/allocation_test.jl
index e720c042d..38bd576d3 100644
--- a/core/test/allocation_test.jl
+++ b/core/test/allocation_test.jl
@@ -44,7 +44,7 @@
     @test allocated[3, :] ≈ [0.0, 2.0]
 end
 
-@testitem "Allocation objective: linear absolute" begin
+@testitem "Allocation objective" begin
     using DataFrames: DataFrame
     using SciMLBase: successful_retcode
     using Ribasim: NodeID
@@ -67,13 +67,13 @@ end
     to_user_5 = F[(NodeID(:Basin, 4), NodeID(:UserDemand, 5))]
     to_user_6 = F[(NodeID(:Basin, 4), NodeID(:UserDemand, 6))]
 
-    @test objective.aff.constant == 2.0
-    @test objective.aff.terms[to_user_5] ≈ -2.0 / user_demand.demand[1]
-    @test objective.aff.terms[to_user_6] ≈ -2.0 / user_demand.demand[2]
+    @test objective.aff.constant ≈ sum(user_demand.demand)
+    @test objective.aff.terms[to_user_5] ≈ -2.0
+    @test objective.aff.terms[to_user_6] ≈ -2.0
     @test objective.terms[JuMP.UnorderedPair(to_user_5, to_user_5)] ≈
-          1 / user_demand.demand[1]^2
+          1 / user_demand.demand[1]
     @test objective.terms[JuMP.UnorderedPair(to_user_6, to_user_6)] ≈
-          1 / user_demand.demand[2]^2
+          1 / user_demand.demand[2]
 end
 
 @testitem "Allocation with controlled fractional flow" begin
@@ -168,7 +168,7 @@ end
     @test Ribasim.is_main_network(first(allocation_network_ids))
 
     # Connections from main network to subnetworks
-    @test isempty(main_network_connections[1])
+    @test only(main_network_connections[1]) == (NodeID(:FlowBoundary, 1), NodeID(:Basin, 2))
     @test only(main_network_connections[2]) == (NodeID(:Basin, 2), NodeID(:Pump, 11))
     @test only(main_network_connections[3]) == (NodeID(:Basin, 6), NodeID(:Pump, 24))
     @test only(main_network_connections[4]) == (NodeID(:Basin, 10), NodeID(:Pump, 38))
@@ -263,7 +263,7 @@ end
     @test all(
         isapprox.(
             subnetwork_allocateds[NodeID(:Basin, 2), NodeID(:Pump, 11)],
-            [4.0, 0.493, 0.0],
+            [4.0, 0.4947, 0.0],
             atol = 1e-4,
         ),
     )
@@ -275,7 +275,7 @@ end
     @test all(
         isapprox.(
             subnetwork_allocateds[NodeID(:Basin, 10), NodeID(:Pump, 38)],
-            [0.001, 0.002, 0.0],
+            [0.001, 2.473e-4, 0.0],
             rtol = 1e-3,
         ),
     )
@@ -435,7 +435,6 @@ end
     F = problem[:F]
     F_flow_buffer_in = problem[:F_flow_buffer_in]
     F_flow_buffer_out = problem[:F_flow_buffer_out]
-    F_abs_flow_demand = problem[:F_abs_flow_demand]
 
     node_id_with_flow_demand = NodeID(NodeType.TabulatedRatingCurve, 2)
     constraint_flow_out = problem[:flow_demand_outflow][node_id_with_flow_demand]
@@ -472,7 +471,7 @@ end
     )
     objective = JuMP.objective_function(problem)
     # Reduced demand
-    @test isapprox(flow_demand.demand[1], flow_demand.demand_itp[1](t) - 0.001, rtol = 1e-4)
+    @test isapprox(flow_demand.demand[1], flow_demand.demand_itp[1](t) - 0.001, rtol = 1e-3)
     @test JuMP.normalized_rhs(constraint_flow_out) == Inf
 
     ## Priority 2
@@ -487,7 +486,7 @@ end
     # No demand left
     @test flow_demand.demand[1] < 1e-10
     # Allocated
-    @test isapprox(JuMP.value(only(F_flow_buffer_in)), 0.001, rtol = 1e-4)
+    @test isapprox(JuMP.value(only(F_flow_buffer_in)), 0.001, rtol = 1e-3)
     @test JuMP.normalized_rhs(constraint_flow_out) == 0.0
 
     ## Priority 3
@@ -506,13 +505,13 @@ end
     @test isapprox(
         JuMP.value(F_flow_buffer_out[node_id_with_flow_demand]),
         0.001,
-        rtol = 1e-4,
+        rtol = 1e-3,
     )
     # Flow taken from buffer
     @test isapprox(
         JuMP.value(only(F_flow_buffer_out)),
         user_demand.demand_itp[1][3](t),
-        rtol = 1e-4,
+        rtol = 1e-3,
     )
     # No flow coming from level boundary
     @test JuMP.value(F[(only(level_boundary.node_id), node_id_with_flow_demand)]) == 0
@@ -528,11 +527,11 @@ end
     )
     # Get demand from buffers
     d = user_demand.demand_itp[3][4](t)
-    @assert isapprox(
+    @test isapprox(
         JuMP.value(F[(NodeID(NodeType.UserDemand, 4), NodeID(NodeType.Basin, 7))]) +
         JuMP.value(F[(NodeID(NodeType.UserDemand, 6), NodeID(NodeType.Basin, 7))]),
         d,
-        rtol = 1e-4,
+        rtol = 1e-3,
     )
 end
 
@@ -557,3 +556,29 @@ end
     )
     @test constraint.set.upper == 2.0
 end
+
+@testitem "equal_fraction_allocation" begin
+    using Ribasim: NodeID, NodeType
+    using StructArrays: StructVector
+    using DataFrames: DataFrame
+
+    toml_path =
+        normpath(@__DIR__, "../../generated_testmodels/fair_distribution/ribasim.toml")
+    @test ispath(toml_path)
+    model = Ribasim.run(toml_path)
+    (; user_demand, graph) = model.integrator.p
+
+    data_allocation = DataFrame(Ribasim.allocation_table(model))
+    fractions = Vector{Float64}[]
+
+    for id in user_demand.node_id
+        data_allocation_id =
+            filter(:node_id => node_id -> node_id == id.value, data_allocation)
+        frac = data_allocation_id.allocated ./ data_allocation_id.demand
+        push!(fractions, frac)
+    end
+
+    @test all(isapprox.(fractions[1], fractions[2], atol = 1e-4))
+    @test all(isapprox.(fractions[1], fractions[3], atol = 1e-4))
+    @test all(isapprox.(fractions[1], fractions[4], atol = 1e-4))
+end
diff --git a/python/ribasim_testmodels/ribasim_testmodels/allocation.py b/python/ribasim_testmodels/ribasim_testmodels/allocation.py
index df2f8383e..67ac25ce3 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/allocation.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/allocation.py
@@ -1039,23 +1039,22 @@ def fair_distribution_model():
         allocation=Allocation(use_allocation=True),
     )
 
-    model.flow_boundary.add(
+    model.level_boundary.add(
         Node(1, Point(0, 0), subnetwork_id=1),
         [
-            flow_boundary.Time(
-                time=pd.date_range(start="2020-01", end="2020-05", freq="MS"),
-                flow_rate=np.arange(0, 12.5, 2.5),
+            level_boundary.Static(
+                level=[1.0],
             )
         ],
     )
 
-    model.linear_resistance.add(
+    model.pump.add(
         Node(
             2,
             Point(1, 0),
             subnetwork_id=1,
         ),
-        [linear_resistance.Static(flow_rate=[9.0], max_flow_rate=9.0)],
+        [pump.Static(flow_rate=9.0, max_flow_rate=[9.0])],
     )
 
     model.basin.add(
@@ -1086,7 +1085,7 @@ def fair_distribution_model():
         Node(7, Point(2, -1), subnetwork_id=1),
         [
             user_demand.Static(
-                priority=[1], demand=1.0, return_factor=1.0, min_level=0.2
+                priority=[1], demand=2.0, return_factor=1.0, min_level=0.2
             )
         ],
     )
@@ -1095,7 +1094,7 @@ def fair_distribution_model():
         Node(8, Point(4, 1), subnetwork_id=1),
         [
             user_demand.Static(
-                priority=[1], demand=1.0, return_factor=1.0, min_level=0.2
+                priority=[1], demand=3.0, return_factor=1.0, min_level=0.2
             )
         ],
     )
@@ -1103,13 +1102,17 @@ def fair_distribution_model():
     model.user_demand.add(
         Node(9, Point(4, -1), subnetwork_id=1),
         [
-            user_demand.Static(
-                priority=[1], demand=1.0, return_factor=1.0, min_level=0.2
+            user_demand.Time(
+                priority=1,
+                time=pd.date_range(start="2020-01", end="2021-01", freq="MS"),
+                demand=np.linspace(1.0, 5.0, 13),
+                return_factor=1.0,
+                min_level=0.2,
             )
         ],
     )
 
-    model.edge.add(model.flow_boundary[1], model.pump[2])
+    model.edge.add(model.level_boundary[1], model.pump[2], subnetwork_id=1)
     model.edge.add(model.pump[2], model.basin[3])
     model.edge.add(model.basin[3], model.linear_resistance[4])
     model.edge.add(model.linear_resistance[4], model.basin[5])
@@ -1117,5 +1120,9 @@ def fair_distribution_model():
     model.edge.add(model.basin[3], model.user_demand[7])
     model.edge.add(model.basin[5], model.user_demand[8])
     model.edge.add(model.basin[5], model.user_demand[9])
+    model.edge.add(model.user_demand[6], model.basin[3])
+    model.edge.add(model.user_demand[7], model.basin[3])
+    model.edge.add(model.user_demand[8], model.basin[5])
+    model.edge.add(model.user_demand[9], model.basin[5])
 
     return model

From 1961d5721cfc6b94eb33c26578866caf1990c3c0 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Wed, 10 Apr 2024 08:55:25 +0200
Subject: [PATCH 6/7] Update docs

---
 docs/core/allocation.qmd | 12 +++++-------
 1 file changed, 5 insertions(+), 7 deletions(-)

diff --git a/docs/core/allocation.qmd b/docs/core/allocation.qmd
index e3669da7d..b0aa5e336 100644
--- a/docs/core/allocation.qmd
+++ b/docs/core/allocation.qmd
@@ -158,25 +158,23 @@ $$
 
 The error between the flows and user demands is denoted by $E_{\text{user demand}}$, where
 $$
-    E_{\text{user demand}} = \sum_{(i,j)\in E_S\;:\; i\in U_S} \left| F_{ij} - d_j^p(t)\right|
+    E_{\text{user demand}} = \sum_{(i,j)\in E_S\;:\; i\in U_S} d_j^p(t)\left(1 - \frac{F_{ij}}{d_j^p(t)}\right)^2
 $$
 
 :::{.callout-note}
 When performing main network allocation, the connections to subnetworks are also interpreted as UserDemand with demands determined by subnetwork demand collection.
 :::
 
-This type of objective cares about the absolute amount of water allocated to a demand. It treats all deviations equally which means it doesn't give larger punishment per flow unit if deviations increase.
-
-The absolute value applied here is not supported in a linear programming context directly; this requires introduction of new variables and constraints. For more details see [here](https://optimization.cbe.cornell.edu/index.php?title=Optimization_with_absolute_values).
+This type of objective cares about the fraction of the demand allocated, and will lead to an equal fraction of all demands allocated when possible.
 
 Likewise, the error of level demands from basins is the absolute difference between flows consumed by basins and basin demands.
 $$
-    E_{\text{level demand}} = \sum_{i \in B_S} \left| F_i^\text{basin in} - d_i^p(t)\right|
+    E_{\text{level demand}} = \sum_{i \in B_S} d_i^p(t)\left(1 - \frac{F_i^\text{basin in}}{d_i^p(t)}\right)^2
 $$
 
 Lastly, the error of the flow demands is given as below.
 $$
-    E_{\text{flow demand}} = \sum_{i \in FD_S} \left| F_i^\text{buffer in} - d_i^p(t)\right|
+    E_{\text{flow demand}} = \sum_{i \in FD_S} d_i^p(t)\left(1 -  \frac{F_i^\text{buffer in}}{d_i^p(t)}\right)^2
 $$
 
 ## The optimization constraints
@@ -238,7 +236,7 @@ $$ {#eq-fractionalflowconstraint}
 
 ## Example
 
-The following is an example of an optimization problem for the example shown [here](../python/examples.ipynb#model-with-allocation-user-demand):
+The following is an example of an optimization problem for the example model shown [here](../python/examples.ipynb#model-with-allocation-user-demand):
 
 
 ```{julia}

From ab91921bc6432dab685ebc40be5615a712d16f82 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Thu, 11 Apr 2024 16:51:01 +0200
Subject: [PATCH 7/7] Comments adressed

---
 core/src/allocation_optim.jl                  |  6 +-
 core/test/allocation_test.jl                  | 83 +++++--------------
 .../ribasim_testmodels/allocation.py          |  2 +-
 3 files changed, 26 insertions(+), 65 deletions(-)

diff --git a/core/src/allocation_optim.jl b/core/src/allocation_optim.jl
index 07bc1ff29..9172c13b2 100644
--- a/core/src/allocation_optim.jl
+++ b/core/src/allocation_optim.jl
@@ -1,7 +1,7 @@
 @enumx OptimizationType internal_sources collect_demands allocate
 
 """
-Add an objective term (1 - flow/demand)^2. If the absolute
+Add an objective term `demand * (1 - flow/demand)^2`. If the absolute
 value of the demand is very small, this would lead to huge coefficients,
 so in that case a term of the form (flow - demand)^2 is used.
 """
@@ -11,12 +11,12 @@ function add_objective_term!(
     F::JuMP.VariableRef,
 )::Nothing
     if abs(demand) < 1e-5
-        # Error term (F - d)^2
+        # Error term (F - d)^2 = F² - 2dF + d²
         JuMP.add_to_expression!(ex, 1.0, F, F)
         JuMP.add_to_expression!(ex, -2.0 * demand, F)
         JuMP.add_to_expression!(ex, demand^2)
     else
-        # Error term d*(1 - F/d)^2
+        # Error term d*(1 - F/d)^2 = F²/d - 2F + d
         JuMP.add_to_expression!(ex, 1.0 / demand, F, F)
         JuMP.add_to_expression!(ex, -2.0, F)
         JuMP.add_to_expression!(ex, demand)
diff --git a/core/test/allocation_test.jl b/core/test/allocation_test.jl
index 38bd576d3..262c4e736 100644
--- a/core/test/allocation_test.jl
+++ b/core/test/allocation_test.jl
@@ -223,21 +223,11 @@ end
 
     # See the difference between these values here and in
     # "subnetworks_with_sources"
-    @test all(
-        isapprox.(
-            subnetwork_demands[(NodeID(:Basin, 2), NodeID(:Pump, 11))],
-            [4.0, 4.0, 0.0],
-            atol = 1e-4,
-        ),
-    )
+    @test subnetwork_demands[(NodeID(:Basin, 2), NodeID(:Pump, 11))] ≈ [4.0, 4.0, 0.0] atol =
+        1e-4
     @test subnetwork_demands[(NodeID(:Basin, 6), NodeID(:Pump, 24))] ≈ [0.004, 0.0, 0.0]
-    @test all(
-        isapprox.(
-            subnetwork_demands[(NodeID(:Basin, 10), NodeID(:Pump, 38))][1:2],
-            [0.001, 0.002],
-            rtol = 1e-4,
-        ),
-    )
+    @test subnetwork_demands[(NodeID(:Basin, 10), NodeID(:Pump, 38))][1:2] ≈ [0.001, 0.002] rtol =
+        1e-4
 
     # Solving for the main network, containing subnetworks as UserDemands
     allocation_model = allocation_models[1]
@@ -260,25 +250,13 @@ end
     u = ComponentVector(; storage = zeros(length(p.basin.node_id)))
     Ribasim.update_allocation!((; p, t, u))
 
-    @test all(
-        isapprox.(
-            subnetwork_allocateds[NodeID(:Basin, 2), NodeID(:Pump, 11)],
-            [4.0, 0.4947, 0.0],
-            atol = 1e-4,
-        ),
-    )
-    @test isapprox(
-        subnetwork_allocateds[NodeID(:Basin, 6), NodeID(:Pump, 24)],
-        [0.004, 0.0, 0.0],
-        rtol = 1e-3,
-    )
-    @test all(
-        isapprox.(
-            subnetwork_allocateds[NodeID(:Basin, 10), NodeID(:Pump, 38)],
-            [0.001, 2.473e-4, 0.0],
-            rtol = 1e-3,
-        ),
-    )
+    @test subnetwork_allocateds[NodeID(:Basin, 2), NodeID(:Pump, 11)] ≈ [4.0, 0.4947, 0.0] atol =
+        1e-4
+    @test subnetwork_allocateds[NodeID(:Basin, 6), NodeID(:Pump, 24)] ≈ [0.004, 0.0, 0.0] rtol =
+        1e-3
+
+    @test subnetwork_allocateds[NodeID(:Basin, 10), NodeID(:Pump, 38)] ≈
+          [0.001, 2.473e-4, 0.0] rtol = 1e-3
 
     # Test for existence of edges in allocation flow record
     allocation_flow = DataFrame(record_flow)
@@ -322,7 +300,6 @@ end
     # Collecting demands
     u = ComponentVector(; storage = zeros(length(basin.node_id)))
     for (i, allocation_model) in enumerate(allocation_models[2:end])
-        @show i
         Ribasim.allocate!(p, allocation_model, t, u, OptimizationType.internal_sources)
         Ribasim.allocate!(p, allocation_model, t, u, OptimizationType.collect_demands)
     end
@@ -330,13 +307,8 @@ end
     # See the difference between these values here and in
     # "allocation with main network optimization problem", internal sources
     # lower the subnetwork demands
-    @test all(
-        isapprox.(
-            subnetwork_demands[(NodeID(:Basin, 2), NodeID(:Pump, 11))],
-            [3.1, 4.0, 0.0],
-            atol = 1e-4,
-        ),
-    )
+    @test subnetwork_demands[(NodeID(:Basin, 2), NodeID(:Pump, 11))] ≈ [3.1, 4.0, 0.0] atol =
+        1e-4
     @test subnetwork_demands[(NodeID(:Basin, 6), NodeID(:Pump, 24))] ≈ [0.004, 0.0, 0.0]
     @test subnetwork_demands[(NodeID(:Basin, 10), NodeID(:Pump, 38))][1:2] ≈ [0.001, 0.001]
 end
@@ -471,7 +443,7 @@ end
     )
     objective = JuMP.objective_function(problem)
     # Reduced demand
-    @test isapprox(flow_demand.demand[1], flow_demand.demand_itp[1](t) - 0.001, rtol = 1e-3)
+    @test flow_demand.demand[1] ≈ flow_demand.demand_itp[1](t) - 0.001 rtol = 1e-3
     @test JuMP.normalized_rhs(constraint_flow_out) == Inf
 
     ## Priority 2
@@ -486,7 +458,7 @@ end
     # No demand left
     @test flow_demand.demand[1] < 1e-10
     # Allocated
-    @test isapprox(JuMP.value(only(F_flow_buffer_in)), 0.001, rtol = 1e-3)
+    @test JuMP.value(only(F_flow_buffer_in)) ≈ 0.001 rtol = 1e-3
     @test JuMP.normalized_rhs(constraint_flow_out) == 0.0
 
     ## Priority 3
@@ -502,17 +474,10 @@ end
     # The flow from the source is used up in previous priorities
     @test JuMP.value(F[(NodeID(NodeType.LevelBoundary, 1), node_id_with_flow_demand)]) == 0
     # So flow from the flow buffer is used for UserDemand #4
-    @test isapprox(
-        JuMP.value(F_flow_buffer_out[node_id_with_flow_demand]),
-        0.001,
-        rtol = 1e-3,
-    )
+    @test JuMP.value(F_flow_buffer_out[node_id_with_flow_demand]) ≈ 0.001 rtol = 1e-3
+
     # Flow taken from buffer
-    @test isapprox(
-        JuMP.value(only(F_flow_buffer_out)),
-        user_demand.demand_itp[1][3](t),
-        rtol = 1e-3,
-    )
+    @test JuMP.value(only(F_flow_buffer_out)) ≈ user_demand.demand_itp[1][3](t) rtol = 1e-3
     # No flow coming from level boundary
     @test JuMP.value(F[(only(level_boundary.node_id), node_id_with_flow_demand)]) == 0
 
@@ -527,12 +492,9 @@ end
     )
     # Get demand from buffers
     d = user_demand.demand_itp[3][4](t)
-    @test isapprox(
-        JuMP.value(F[(NodeID(NodeType.UserDemand, 4), NodeID(NodeType.Basin, 7))]) +
-        JuMP.value(F[(NodeID(NodeType.UserDemand, 6), NodeID(NodeType.Basin, 7))]),
-        d,
-        rtol = 1e-3,
-    )
+    @test JuMP.value(F[(NodeID(NodeType.UserDemand, 4), NodeID(NodeType.Basin, 7))]) +
+          JuMP.value(F[(NodeID(NodeType.UserDemand, 6), NodeID(NodeType.Basin, 7))]) ≈ d rtol =
+        1e-3
 end
 
 @testitem "flow_demand_with_max_flow_rate" begin
@@ -572,8 +534,7 @@ end
     fractions = Vector{Float64}[]
 
     for id in user_demand.node_id
-        data_allocation_id =
-            filter(:node_id => node_id -> node_id == id.value, data_allocation)
+        data_allocation_id = filter(:node_id => ==(id.value), data_allocation)
         frac = data_allocation_id.allocated ./ data_allocation_id.demand
         push!(fractions, frac)
     end
diff --git a/python/ribasim_testmodels/ribasim_testmodels/allocation.py b/python/ribasim_testmodels/ribasim_testmodels/allocation.py
index 67ac25ce3..5172a8dd8 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/allocation.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/allocation.py
@@ -1034,7 +1034,7 @@ def linear_resistance_demand_model():
 def fair_distribution_model():
     model = Model(
         starttime="2020-01-01 00:00:00",
-        endtime="2021-01-01 00:00:00",
+        endtime="2020-01-07 00:00:00",
         crs="EPSG:28992",
         allocation=Allocation(use_allocation=True),
     )