Add integration for allocation

core/src/allocation_optim.jl

@@ -255,7 +255,7 @@ function set_initial_capacities_source!(
 )::Nothing
     (; problem) = allocation_model
     (; graph, allocation) = p
-    (; mean_flows) = allocation
+    (; integrated_flow, integrated_flow_mapping) = allocation
     (; subnetwork_id) = allocation_model
     source_constraints = problem[:source]
     main_network_source_edges = get_main_network_connections(p, subnetwork_id)
@@ -266,7 +266,7 @@ function set_initial_capacities_source!(
             # If it is a source edge for this allocation problem
             if edge ∉ main_network_source_edges
                 # Reset the source to the averaged flow over the last allocation period
-                source_capacity = mean_flows[edge][]
+                source_capacity = integrated_flow.integrand[integrated_flow_mapping[edge]]
                 JuMP.set_normalized_rhs(
                     source_constraints[edge],
                     # It is assumed that the allocation procedure does not have to be differentiated.
@@ -361,11 +361,11 @@ function get_basin_data(
     (; graph, basin, level_demand, allocation) = p
     (; vertical_flux) = basin
     (; Δt_allocation) = allocation_model
-    (; mean_flows) = allocation
+    (; integrated_flow, integrated_flow_mapping) = allocation
     @assert node_id.type == NodeType.Basin
     vertical_flux = get_tmp(vertical_flux, 0)
     _, basin_idx = id_index(basin.node_id, node_id)
-    influx = mean_flows[(node_id, node_id)][]
+    influx = integrated_flow.integrand[integrated_flow_mapping[node_id, node_id]]
     _, basin_idx = id_index(basin.node_id, node_id)
     storage_basin = u.storage[basin_idx]
     control_inneighbors = inneighbor_labels_type(graph, node_id, EdgeType.control)

core/src/callback.jl

@@ -9,7 +9,14 @@ function create_callbacks(
     config::Config,
     saveat,
 )::Tuple{CallbackSet, SavedResults}
-    (; starttime, basin, user_demand, tabulated_rating_curve, discrete_control) = parameters
+    (;
+        starttime,
+        basin,
+        user_demand,
+        tabulated_rating_curve,
+        discrete_control,
+        allocation,
+    ) = parameters
     callbacks = SciMLBase.DECallback[]
 
     # Check for negative storages
@@ -20,23 +27,27 @@ function create_callbacks(
     integrand_output_prototype = get_flow_integrand_output_prototype(parameters)
     saved_flow = IntegrandValues(Float64, typeof(integrand_output_prototype))
     integrating_flow_cb =
-        IntegratingCallback(flow_integrand, saved_flow, integrand_output_prototype)
+        IntegratingCallback(flow_integrand!, saved_flow, integrand_output_prototype)
     push!(callbacks, integrating_flow_cb)
 
     # Integrate vertical basin fluxes for BMI
-    stored_for_bmi = IntegrandValuesSum(
-        ComponentVector{Float64}(;
-            get_tmp(basin.vertical_flux, 0)...,
-            user_realized = zeros(length(user_demand.node_id)),
-        ),
-    )
-    integrating_for_bmi_cb = IntegratingSumCallback(
-        bmi_integrand,
-        stored_for_bmi,
-        get_tmp(basin.vertical_flux, 0),
+    integrand_output_prototype = ComponentVector{Float64}(;
+        get_tmp(basin.vertical_flux, 0)...,
+        user_realized = zeros(length(user_demand.node_id)),
     )
+    stored_for_bmi = IntegrandValuesSum(integrand_output_prototype)
+    integrating_for_bmi_cb =
+        IntegratingSumCallback(bmi_integrand!, stored_for_bmi, integrand_output_prototype)
     push!(callbacks, integrating_for_bmi_cb)
 
+    # Integrate flows for allocation input
+    integrating_for_allocation_cb = IntegratingSumCallback(
+        allocation_integrand!,
+        allocation.integrated_flow,
+        zeros(length(allocation.integrated_flow_mapping)),
+    )
+    push!(callbacks, integrating_for_allocation_cb)
+
     # TODO: What is this?
     tstops = get_tstops(basin.time.time, starttime)
     basin_cb = PresetTimeCallback(tstops, update_basin; save_positions = (false, false))
@@ -74,38 +85,20 @@ function create_callbacks(
     return callback, saved
 end
 
-function flow_integrand(out, u, t, integrator)::Nothing
+function flow_integrand!(out, u, t, integrator)::Nothing
     (; basin, graph) = integrator.p
     (; vertical_flux) = basin
     vertical_flux = get_tmp(vertical_flux, 0)
     flow = get_tmp(graph[].flow, 0)
-    vertical_flux_view(out) .= vertical_flux
 
     out.flow .= flow
-
-    for (i, basin_id) in enumerate(basin.node_id)
-        for inflow_edge in basin.inflow_edges[i]
-            q = get_flow(graph, inflow_edge, 0)
-            if q > 0
-                out.basin_inflow[i] += q
-            else
-                out.basin_outflow[i] -= q
-            end
-        end
-        for outflow_edge in basin.outflow_edges[i]
-            q = get_flow(graph, outflow_edge, 0)
-            if q > 0
-                out.basin_outflow[i] += q
-            else
-                out.basin_inflow[i] -= q
-            end
-        end
-    end
+    vertical_flux_view(out) .= vertical_flux
+    set_inoutflows!(out, graph, basin)
     return nothing
 end
 
-function bmi_integrand(out, u, t, integrator)::Nothing
-    (; basin, user_demand) = integrator.p
+function bmi_integrand!(out, u, t, integrator)::Nothing
+    (; basin, user_demand, graph) = integrator.p
     vertical_flux_view(out) .= get_tmp(basin.vertical_flux, 0)
 
     for i in eachindex(user_demand.node_id)
@@ -114,6 +107,21 @@ function bmi_integrand(out, u, t, integrator)::Nothing
     return
 end
 
+function allocation_integrand!(out, u, t, integrator)::Nothing
+    (; allocation, graph, basin) = integrator.p
+    for (edge, i) in allocation.integrated_flow_mapping
+        if edge[1] == edge[2]
+            # Vertical fluxes
+            _, basin_idx = id_index(basin.node_id, edge[1])
+            out[i] = get_influx(basin, basin_idx)
+        else
+            # Horizontal flows
+            out[i] = get_flow(graph, edge..., 0)
+        end
+    end
+    return nothing
+end
+
 function get_flow_integrand_output_prototype(p::Parameters)
     (; graph, basin) = p
     flow = get_tmp(graph[].flow, 0)
@@ -385,7 +393,7 @@ function update_basin(integrator)::Nothing
     (; p, u) = integrator
     (; basin) = p
     (; storage) = u
-    (; node_id, time, vertical_flux_from_input, vertical_flux, vertical_flux_prev) = basin
+    (; node_id, time, vertical_flux_from_input, vertical_flux) = basin
     t = datetime_since(integrator.t, integrator.p.starttime)
     vertical_flux = get_tmp(vertical_flux, integrator.u)
 
@@ -406,17 +414,14 @@ function update_basin(integrator)::Nothing
     end
 
     update_vertical_flux!(basin, storage)
-
-    # Forget about vertical fluxes to handle discontinuous forcing from basin_update
-    copyto!(vertical_flux_prev, vertical_flux)
     return nothing
 end
 
 "Solve the allocation problem for all demands and assign allocated abstractions."
 function update_allocation!(integrator)::Nothing
     (; p, t, u) = integrator
     (; allocation) = p
-    (; allocation_models, mean_flows) = allocation
+    (; allocation_models, integrated_flow) = allocation
 
     # Don't run the allocation algorithm if allocation is not active
     # (Specifically for running Ribasim via the BMI)
@@ -428,9 +433,7 @@ function update_allocation!(integrator)::Nothing
 
     # Divide by the allocation Δt to obtain the mean flows
     # from the integrated flows
-    for value in values(mean_flows)
-        value[] /= Δt_allocation
-    end
+    integrated_flow.integrand ./= Δt_allocation
 
     # If a main network is present, collect demands of subnetworks
     if has_main_network(allocation)
@@ -448,9 +451,8 @@ function update_allocation!(integrator)::Nothing
     end
 
     # Reset the mean source flows
-    for value in values(mean_flows)
-        value[] = 0.0
-    end
+    integrated_flow.integrand .= 0.0
+    return nothing
 end
 
 "Load updates from 'TabulatedRatingCurve / time' into the parameters"

core/src/parameter.jl

@@ -75,7 +75,8 @@ struct Allocation
     priorities::Vector{Int32}
     subnetwork_demands::Dict{Tuple{NodeID, NodeID}, Vector{Float64}}
     subnetwork_allocateds::Dict{Tuple{NodeID, NodeID}, Vector{Float64}}
-    mean_flows::Dict{Tuple{NodeID, NodeID}, Base.RefValue{Float64}}
+    integrated_flow_mapping::Dict{Tuple{NodeID, NodeID}, Int32}
+    integrated_flow::IntegrandValuesSum{Vector{Float64}}
     record_demand::@NamedTuple{
         time::Vector{Float64},
         subnetwork_id::Vector{Int32},

core/src/read.jl

@@ -1018,22 +1018,26 @@ function Allocation(db::DB, config::Config, graph::MetaGraph)::Allocation
         optimization_type = String[],
     )
 
-    mean_flows = Dict{Tuple{NodeID, NodeID}, Base.RefValue{Float64}}()
+    integrated_flow_mapping = Dict{Tuple{NodeID, NodeID}, Int32}()
 
     # Find edges which serve as sources in allocation
+    flow_counter = 0
     for edge_metadata in values(graph.edge_data)
         (; subnetwork_id_source, edge) = edge_metadata
         if subnetwork_id_source != 0
-            mean_flows[edge] = Ref(0.0)
+            flow_counter += 1
+            integrated_flow_mapping[edge] = flow_counter
         end
     end
 
     # Find basins with a level demand
     for node_id in values(graph.vertex_labels)
         if has_external_demand(graph, node_id, :level_demand)[1]
-            mean_flows[(node_id, node_id)] = Ref(0.0)
+            flow_counter += 1
+            integrated_flow_mapping[(node_id, node_id)] = flow_counter
         end
     end
+    integrated_flow = IntegrandValuesSum(zeros(flow_counter))
 
     return Allocation(
         Int32[],
@@ -1042,7 +1046,8 @@ function Allocation(db::DB, config::Config, graph::MetaGraph)::Allocation
         get_all_priorities(db, config),
         Dict{Tuple{NodeID, NodeID}, Vector{Float64}}(),
         Dict{Tuple{NodeID, NodeID}, Vector{Float64}}(),
-        mean_flows,
+        integrated_flow_mapping,
+        integrated_flow,
         record_demand,
         record_flow,
     )

core/src/util.jl

@@ -757,3 +757,29 @@ end
 function vertical_flux_view(v::ComponentVector)
     return @view v[(:precipitation, :evaporation, :drainage, :infiltration)]
 end
+
+function set_inoutflows!(
+    v::ComponentVector{Float64},
+    graph::MetaGraph,
+    basin::Basin,
+)::Nothing
+    for (i, basin_id) in enumerate(basin.node_id)
+        for inflow_edge in basin.inflow_edges[i]
+            q = get_flow(graph, inflow_edge, 0)
+            if q > 0
+                v.basin_inflow[i] += q
+            else
+                v.basin_outflow[i] -= q
+            end
+        end
+        for outflow_edge in basin.outflow_edges[i]
+            q = get_flow(graph, outflow_edge, 0)
+            if q > 0
+                v.basin_outflow[i] += q
+            else
+                v.basin_inflow[i] -= q
+            end
+        end
+    end
+    return nothing
+end

core/src/write.jl

@@ -83,11 +83,12 @@ function average_over_saveats(
 )
     (; integrand, ts) = integrand_values
     averages = Vector{Float64}[]
+    n_saveats = length(saveats)
     saveat_index = 2
     integral_saveat = zero(first(integrand)[symb])
     for (integral_step, t) in zip(integrand, ts)
         integral_saveat += integral_step[symb]
-        if t == saveats[saveat_index]
+        if saveat_index <= n_saveats && t == saveats[saveat_index]
             Δt_saveat = saveats[saveat_index] - saveats[saveat_index - 1]
             push!(averages, copy(integral_saveat) / Δt_saveat)
             integral_saveat .= 0.0