make demand, allocated matrices and index these directly (#1353)

Fixes #1277

core/src/allocation_optim.jl

@@ -78,7 +78,7 @@ function set_objective_priority!(
 )::Nothing
     (; problem, allocation_network_id) = allocation_model
     (; graph, user_demand, flow_demand, allocation, basin) = p
-    (; node_id) = user_demand
+    (; node_id, demand_reduced) = user_demand
     (; main_network_connections, subnetwork_demands) = allocation
     edge_ids = graph[].edge_ids[allocation_network_id]
 
@@ -114,7 +114,8 @@ function set_objective_priority!(
 
         if to_node_id.type == NodeType.UserDemand
             # UserDemand
-            d = get_user_demand(p, to_node_id, priority_idx)
+            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)
         else
             has_demand, demand_node_id =
@@ -183,7 +184,7 @@ function assign_allocations!(
             if user_demand_node_id.type == NodeType.UserDemand
                 allocated = JuMP.value(F[edge_id])
                 user_demand_idx = findsorted(user_demand.node_id, user_demand_node_id)
-                user_demand.allocated[user_demand_idx][priority_idx] = allocated
+                user_demand.allocated[user_demand_idx, priority_idx] = allocated
             end
         end
     end
@@ -481,20 +482,19 @@ function set_initial_demands_user!(
 )::Nothing
     (; allocation_network_id) = allocation_model
     (; graph, user_demand, allocation) = p
-    (; node_id, demand_from_timeseries, demand_itp) = user_demand
+    (; node_id, demand_from_timeseries, demand_itp, demand, demand_reduced) = user_demand
 
     # Read the demand from the interpolated timeseries
     # for users for which the demand comes from there
     for (i, id) in enumerate(node_id)
         if demand_from_timeseries[i] &&
            graph[id].allocation_network_id == allocation_network_id
             for priority_idx in eachindex(allocation.priorities)
-                d = demand_itp[i][priority_idx](t)
-                set_user_demand!(p, id, priority_idx, d; reduced = false)
+                demand[i, priority_idx] = demand_itp[i][priority_idx](t)
             end
         end
     end
-    copy!(user_demand.demand_reduced, user_demand.demand)
+    copy!(demand_reduced, demand)
     return nothing
 end
 
@@ -576,17 +576,17 @@ function adjust_demands_user!(
 )::Nothing
     (; problem, allocation_network_id) = allocation_model
     (; graph, user_demand) = p
+    (; node_id, demand_reduced) = user_demand
     F = problem[:F]
 
     # Reduce the demand by what was allocated
-    for id in user_demand.node_id
+    for (i, id) in enumerate(node_id)
         if graph[id].allocation_network_id == allocation_network_id
             d = max(
                 0.0,
-                get_user_demand(p, id, priority_idx) -
-                JuMP.value(F[(inflow_id(graph, id), id)]),
+                demand_reduced[i, priority_idx] - JuMP.value(F[(inflow_id(graph, id), id)]),
             )
-            set_user_demand!(p, id, priority_idx, d)
+            demand_reduced[i, priority_idx] = d
         end
     end
     return nothing
@@ -755,8 +755,8 @@ function save_demands_and_allocations!(
         if node_id.type == NodeType.UserDemand
             has_demand = true
             user_demand_idx = findsorted(user_demand.node_id, node_id)
-            demand = get_user_demand(p, node_id, priority_idx; reduced = false)
-            allocated = user_demand.allocated[user_demand_idx][priority_idx]
+            demand = user_demand.demand[user_demand_idx, priority_idx]
+            allocated = user_demand.allocated[user_demand_idx, priority_idx]
             #NOTE: instantaneous
             realized = get_flow(graph, inflow_id(graph, node_id), node_id, 0)
 

core/src/parameter.jl

@@ -496,11 +496,11 @@ struct UserDemand <: AbstractParameterNode
     node_id::Vector{NodeID}
     active::BitVector
     realized_bmi::Vector{Float64}
-    demand::Vector{Float64}
-    demand_reduced::Vector{Float64}
+    demand::Matrix{Float64}
+    demand_reduced::Matrix{Float64}
     demand_itp::Vector{Vector{ScalarInterpolation}}
     demand_from_timeseries::BitVector
-    allocated::Vector{Vector{Float64}}
+    allocated::Matrix{Float64}
     return_factor::Vector{Float64}
     min_level::Vector{Float64}
 

core/src/read.jl

@@ -652,7 +652,7 @@ end
 
 function user_demand_static!(
     active::BitVector,
-    demand::Vector{Float64},
+    demand::Matrix{Float64},
     demand_itp::Vector{Vector{ScalarInterpolation}},
     return_factor::Vector{Float64},
     min_level::Vector{Float64},
@@ -672,15 +672,15 @@ function user_demand_static!(
         for row in group
             priority_idx = findsorted(priorities, row.priority)
             demand_itp[user_demand_idx][priority_idx].u .= row.demand
-            demand[(user_demand_idx - 1) * length(priorities) + priority_idx] = row.demand
+            demand[user_demand_idx, priority_idx] = row.demand
         end
     end
     return nothing
 end
 
 function user_demand_time!(
     active::BitVector,
-    demand::Vector{Float64},
+    demand::Matrix{Float64},
     demand_itp::Vector{Vector{ScalarInterpolation}},
     demand_from_timeseries::BitVector,
     return_factor::Vector{Float64},
@@ -712,8 +712,7 @@ function user_demand_time!(
             :demand;
             default_value = 0.0,
         )
-        demand[(user_demand_idx - 1) * length(priorities) + priority_idx] =
-            demand_p_itp(0.0)
+        demand[user_demand_idx, priority_idx] = demand_p_itp(0.0)
 
         if is_valid
             demand_itp[user_demand_idx][priority_idx] = demand_p_itp
@@ -741,14 +740,14 @@ function UserDemand(db::DB, config::Config)::UserDemand
     n_priority = length(priorities)
     active = BitVector(ones(Bool, n_user))
     realized_bmi = zeros(n_user)
-    demand = zeros(n_user * n_priority)
-    demand_reduced = zeros(n_user * n_priority)
+    demand = zeros(n_user, n_priority)
+    demand_reduced = zeros(n_user, n_priority)
     trivial_timespan = [nextfloat(-Inf), prevfloat(Inf)]
     demand_itp = [
         [LinearInterpolation(zeros(2), trivial_timespan) for i in eachindex(priorities)] for j in eachindex(node_ids)
     ]
     demand_from_timeseries = BitVector(zeros(Bool, n_user))
-    allocated = [fill(Inf, length(priorities)) for id in node_ids]
+    allocated = fill(Inf, n_user, n_priority)
     return_factor = zeros(n_user)
     min_level = zeros(n_user)
 

core/src/solve.jl

@@ -266,12 +266,13 @@ function formulate_flow!(
     storage::AbstractVector,
     t::Number,
 )::Nothing
-    (; graph, basin) = p
+    (; graph, basin, allocation) = p
     (;
         node_id,
         allocated,
         active,
         demand_itp,
+        demand,
         return_factor,
         min_level,
         demand_from_timeseries,
@@ -291,12 +292,12 @@ function formulate_flow!(
         # and the current demand.
         # If allocation is not optimized then allocated = Inf, so the result is always
         # effectively allocated = demand.
-        for priority_idx in eachindex(allocated[i])
-            alloc_prio = allocated[i][priority_idx]
+        for priority_idx in eachindex(allocation.priorities)
+            alloc_prio = allocated[i, priority_idx]
             demand_prio = if demand_from_timeseries[i]
                 demand_itp[i][priority_idx](t)
             else
-                get_user_demand(p, id, priority_idx; reduced = false)
+                demand[i, priority_idx]
             end
             alloc = min(alloc_prio, demand_prio)
             q += alloc

core/src/util.jl

@@ -583,50 +583,6 @@ function is_main_network(allocation_network_id::Int32)::Bool
     return allocation_network_id == 1
 end
 
-function get_user_demand(
-    p::Parameters,
-    node_id::NodeID,
-    priority_idx::Int;
-    reduced::Bool = true,
-)::Float64
-    (; user_demand, allocation) = p
-    (; demand, demand_reduced) = user_demand
-    user_demand_idx = findsorted(user_demand.node_id, node_id)
-    n_priorities = length(allocation.priorities)
-    idx = (user_demand_idx - 1) * n_priorities + priority_idx
-
-    if reduced
-        @assert is_active(allocation) "reduced demand accessed in non-allocation run"
-        return demand_reduced[idx]
-    else
-        return demand[idx]
-    end
-end
-
-function set_user_demand!(
-    p::Parameters,
-    node_id::NodeID,
-    priority_idx::Int,
-    value::Float64;
-    reduced::Bool = true,
-)::Nothing
-    (; user_demand, allocation) = p
-    (; demand, demand_reduced) = user_demand
-
-    user_demand_idx = findsorted(user_demand.node_id, node_id)
-    n_priorities = length(allocation.priorities)
-    idx = (user_demand_idx - 1) * n_priorities + priority_idx
-
-    if reduced
-        @assert is_active(allocation) "reduced demand accessed in non-allocation run"
-        demand_reduced[idx] = value
-    else
-        demand[idx] = value
-    end
-
-    return nothing
-end
-
 function get_all_priorities(db::DB, config::Config)::Vector{Int32}
     priorities = Set{Int32}()
 

core/test/allocation_test.jl

@@ -38,16 +38,10 @@
     @test JuMP.value(F[(NodeID(:FlowBoundary, 1), NodeID(:Basin, 2))]) ≈ 0.5
     @test JuMP.value(F[(NodeID(:Basin, 6), NodeID(:UserDemand, 11))]) ≈ 0.0
 
-    allocated = p.user_demand.allocated
-    @test allocated[1] ≈ [0.0, 0.5]
-    @test allocated[2] ≈ [4.0, 0.0]
-    @test allocated[3] ≈ [0.0, 2.0]
-
-    # Test getting and setting UserDemand demands
-    (; user_demand) = p
-    Ribasim.set_user_demand!(p, NodeID(:UserDemand, 11), 2, Float64(π); reduced = false)
-    @test user_demand.demand[4] ≈ π
-    @test Ribasim.get_user_demand(p, NodeID(:UserDemand, 11), 2; reduced = false) ≈ π
+    (; allocated) = p.user_demand
+    @test allocated[1, :] ≈ [0.0, 0.5]
+    @test allocated[2, :] ≈ [4.0, 0.0]
+    @test allocated[3, :] ≈ [0.0, 2.0]
 end
 
 @testitem "Allocation objective: linear absolute" begin
@@ -197,7 +191,7 @@ end
           [(NodeID(:Basin, 10), NodeID(:Pump, 38))]
 end
 
-@testitem "allocation with main network optimization problem" begin
+@testitem "Allocation with main network optimization problem" begin
     using SQLite
     using Ribasim: NodeID, OptimizationType
     using ComponentArrays: ComponentVector
@@ -255,8 +249,8 @@ end
           [0.00399999999, 0.0, 0.0]
     @test subnetwork_allocateds[NodeID(:Basin, 10), NodeID(:Pump, 38)] ≈ [0.001, 0.0, 0.0]
 
-    @test user_demand.allocated[2] ≈ [4.0, 0.0, 0.0]
-    @test user_demand.allocated[7] ≈ [0.001, 0.0, 0.0]
+    @test user_demand.allocated[2, :] ≈ [4.0, 0.0, 0.0]
+    @test user_demand.allocated[7, :] ≈ [0.001, 0.0, 0.0]
 end
 
 @testitem "subnetworks with sources" begin