From a18de961cba156738d746bcd767c092d357921e4 Mon Sep 17 00:00:00 2001
From: Bart de Koning <bart.dekoning@deltares.nl>
Date: Mon, 22 Jan 2024 16:33:52 +0100
Subject: [PATCH] Fix existing tests

---
 Manifest.toml                                             | 2 +-
 Project.toml                                              | 2 ++
 core/src/allocation.jl                                    | 8 +++++---
 core/src/bmi.jl                                           | 2 +-
 core/test/allocation_test.jl                              | 6 +++---
 .../ribasim_testmodels/ribasim_testmodels/allocation.py   | 2 +-
 6 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/Manifest.toml b/Manifest.toml
index a4e91fb27..43112888e 100644
--- a/Manifest.toml
+++ b/Manifest.toml
@@ -2,7 +2,7 @@
 
 julia_version = "1.10.0"
 manifest_format = "2.0"
-project_hash = "a2e54351da64b7dac3109810459579d647b75bb7"
+project_hash = "212d113fb58ab2e43d8fe0283d717923e72a9a88"
 
 [[deps.ADTypes]]
 git-tree-sha1 = "41c37aa88889c171f1300ceac1313c06e891d245"
diff --git a/Project.toml b/Project.toml
index 6df86f2b1..ea6fed355 100644
--- a/Project.toml
+++ b/Project.toml
@@ -7,8 +7,10 @@ Arrow = "69666777-d1a9-59fb-9406-91d4454c9d45"
 BasicModelInterface = "59605e27-edc0-445a-b93d-c09a3a50b330"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 Infiltrator = "5903a43b-9cc3-4c30-8d17-598619ec4e9b"
+JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
 MetaGraphsNext = "fa8bd995-216d-47f1-8a91-f3b68fbeb377"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
+PreallocationTools = "d236fae5-4411-538c-8e31-a6e3d9e00b46"
 ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823"
 Revise = "295af30f-e4ad-537b-8983-00126c2a3abe"
 Ribasim = "aac5e3d9-0b8f-4d4f-8241-b1a7a9632635"
diff --git a/core/src/allocation.jl b/core/src/allocation.jl
index eb178b914..146e5fd7f 100644
--- a/core/src/allocation.jl
+++ b/core/src/allocation.jl
@@ -795,7 +795,7 @@ function add_user_term!(
 
     elseif objective_type == :quadratic_relative
         # Objective function ∑ (1 - F/d)^2
-        if d ≈ 0
+        if demand ≈ 0
             return nothing
         end
         JuMP.add_to_expression!(ex, 1.0 / demand^2, F_edge, F_edge)
@@ -812,16 +812,17 @@ function add_user_term!(
         JuMP.set_normalized_coefficient(
             problem[:abs_positive][node_id_user],
             F_edge,
-            iszero(d) ? 0 : 1 / demand,
+            iszero(demand) ? 0 : 1 / demand,
         )
         JuMP.set_normalized_coefficient(
             problem[:abs_negative][node_id_user],
             F_edge,
-            iszero(d) ? 0 : -1 / demand,
+            iszero(demand) ? 0 : -1 / demand,
         )
     else
         error("Invalid allocation objective type $objective_type.")
     end
+    return nothing
 end
 
 """
@@ -869,6 +870,7 @@ function set_objective_priority!(
 
         user_idx = findsorted(node_id, node_id_user)
         d = demand[user_idx][priority_idx](t)
+        demand_max = max(demand_max, d)
         add_user_term!(ex, edge_id, objective_type, d, allocation_model)
     end
 
diff --git a/core/src/bmi.jl b/core/src/bmi.jl
index ffed0b493..f41b3dd28 100644
--- a/core/src/bmi.jl
+++ b/core/src/bmi.jl
@@ -483,7 +483,7 @@ function get_allocation_model(p::Parameters, allocation_network_id::Int)::Alloca
     (; allocation_network_ids, allocation_models) = allocation
     idx = findsorted(allocation_network_ids, allocation_network_id)
     if isnothing(idx)
-        error("Invalid allocation network id $allocation_network_id.")
+        error("Invalid allocation network ID $allocation_network_id.")
     else
         return allocation_models[idx]
     end
diff --git a/core/test/allocation_test.jl b/core/test/allocation_test.jl
index 10a98c8e6..c18846efe 100644
--- a/core/test/allocation_test.jl
+++ b/core/test/allocation_test.jl
@@ -202,10 +202,10 @@ end
 
     # In each subnetwork, the connection from the main network to the subnetwork is
     # interpreted as a source
-    @test Ribasim.get_allocation_model(p, 2).problem[:source].axes[1] ==
-          Tuple{NodeID, NodeID}[(2, 11)]
     @test Ribasim.get_allocation_model(p, 3).problem[:source].axes[1] ==
+          Tuple{NodeID, NodeID}[(2, 11)]
+    @test Ribasim.get_allocation_model(p, 5).problem[:source].axes[1] ==
           Tuple{NodeID, NodeID}[(6, 24)]
-    @test Ribasim.get_allocation_model(p, 4).problem[:source].axes[1] ==
+    @test Ribasim.get_allocation_model(p, 7).problem[:source].axes[1] ==
           Tuple{NodeID, NodeID}[(10, 38)]
 end
diff --git a/python/ribasim_testmodels/ribasim_testmodels/allocation.py b/python/ribasim_testmodels/ribasim_testmodels/allocation.py
index 0c5b9d1b4..046521092 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/allocation.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/allocation.py
@@ -1214,6 +1214,6 @@ def main_network_with_subnetworks_model():
     # Fix sources
     df = model.network.edge.df
     df.loc[[18, 55, 99], "allocation_network_id"] = None
-    df.loc[128:130, "allocation_network_id"] = [2, 3, 4]
+    df.loc[128:130, "allocation_network_id"] = [3, 5, 7]
 
     return model