Let solver integrate flows

core/src/Ribasim.jl

@@ -24,7 +24,7 @@ import TranscodingStreams
 using Accessors: @set
 using Arrow: Arrow, Table
 using CodecZstd: ZstdCompressor
-using ComponentArrays: ComponentVector
+using ComponentArrays: ComponentVector, ComponentMatrix
 using DataInterpolations: LinearInterpolation, derivative
 using Dates: Dates, DateTime, Millisecond, @dateformat_str
 using DBInterface: execute

core/src/allocation_optim.jl

@@ -252,10 +252,11 @@ as the average flow over the last Δt_allocation of the source in the physical l
 function set_initial_capacities_source!(
     allocation_model::AllocationModel,
     p::Parameters,
+    u::ComponentVector,
 )::Nothing
     (; problem) = allocation_model
     (; graph, allocation) = p
-    (; mean_flows) = allocation
+    (; input_flow_dict) = allocation
     (; subnetwork_id) = allocation_model
     source_constraints = problem[:source]
     main_network_source_edges = get_main_network_connections(p, subnetwork_id)
@@ -266,7 +267,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 = u.flow_allocation_input[input_flow_dict[edge]]
                 JuMP.set_normalized_rhs(
                     source_constraints[edge],
                     # It is assumed that the allocation procedure does not have to be differentiated.
@@ -361,11 +362,11 @@ function get_basin_data(
     (; graph, basin, level_demand, allocation) = p
     (; vertical_flux) = basin
     (; Δt_allocation) = allocation_model
-    (; mean_flows) = allocation
+    (; input_flow_dict) = 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 = u.flow_allocation_input[input_flow_dict[(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)
@@ -573,7 +574,7 @@ end
 """
 Set the demand of the flow demand nodes. 2 cases:
 - Before the first allocation solve, set the demands to their full value;
-- Before an allocation solve, subtract the flow trough the node with a flow demand
+- Before an allocation solve, subtract the flow through the node with a flow demand
   from the total flow demand (which will be used at the priority of the flow demand only).
 """
 function adjust_demands_user!(
@@ -642,7 +643,7 @@ function set_initial_demands_flow!(
 end
 
 """
-Reduce the flow demand based on flow trough the node with the demand.
+Reduce the flow demand based on flow through the node with the demand.
 Flow from any priority counts.
 """
 function adjust_demands_flow!(allocation_model::AllocationModel, p::Parameters)::Nothing
@@ -984,7 +985,7 @@ function set_initial_values!(
     u::ComponentVector,
     t::Float64,
 )::Nothing
-    set_initial_capacities_source!(allocation_model, p)
+    set_initial_capacities_source!(allocation_model, p, u)
     set_initial_capacities_edge!(allocation_model, p)
     set_initial_capacities_basin!(allocation_model, p, u, t)
     set_initial_capacities_buffer!(allocation_model)

core/src/bmi.jl

@@ -40,15 +40,15 @@ function BMI.get_value_ptr(model::Model, name::AbstractString)::AbstractVector{F
     elseif name == "basin.drainage"
         model.integrator.p.basin.vertical_flux_from_input.drainage
     elseif name == "basin.infiltration_integrated"
-        model.integrator.p.basin.vertical_flux_bmi.infiltration
+        model.integrator.u.infiltration_bmi
     elseif name == "basin.drainage_integrated"
-        model.integrator.p.basin.vertical_flux_bmi.drainage
+        model.integrator.u.drainage_bmi
     elseif name == "basin.subgrid_level"
         model.integrator.p.subgrid.level
     elseif name == "user_demand.demand"
         vec(model.integrator.p.user_demand.demand)
     elseif name == "user_demand.realized"
-        model.integrator.p.user_demand.realized_bmi
+        model.integrator.u.realized_user_demand_bmi
     else
         error("Unknown variable $name")
     end

core/src/callback.jl

@@ -8,6 +8,7 @@ Returns the CallbackSet and the SavedValues for flow.
 function create_callbacks(
     parameters::Parameters,
     config::Config,
+    u0::ComponentVector,
     saveat,
 )::Tuple{CallbackSet, SavedResults}
     (; starttime, basin, tabulated_rating_curve, discrete_control) = parameters
@@ -16,9 +17,6 @@ function create_callbacks(
     negative_storage_cb = FunctionCallingCallback(check_negative_storage)
     push!(callbacks, negative_storage_cb)
 
-    integrating_flows_cb = FunctionCallingCallback(integrate_flows!; func_start = false)
-    push!(callbacks, integrating_flows_cb)
-
     tstops = get_tstops(basin.time.time, starttime)
     basin_cb = PresetTimeCallback(tstops, update_basin; save_positions = (false, false))
     push!(callbacks, basin_cb)
@@ -34,7 +32,7 @@ function create_callbacks(
     # If saveat is a vector which contains 0.0 this callback will still be called
     # at t = 0.0 despite save_start = false
     saveat = saveat isa Vector ? filter(x -> x != 0.0, saveat) : saveat
-    saved_vertical_flux = SavedValues(Float64, typeof(basin.vertical_flux_integrated))
+    saved_vertical_flux = SavedValues(Float64, typeof(copy(forcings_integrated(u0))))
     save_vertical_flux_cb =
         SavingCallback(save_vertical_flux, saved_vertical_flux; saveat, save_start = false)
     push!(callbacks, save_vertical_flux_cb)
@@ -86,67 +84,19 @@ function check_negative_storage(u, t, integrator)::Nothing
     return nothing
 end
 
-"""
-Integrate flows over the last timestep
-"""
-function integrate_flows!(u, t, integrator)::Nothing
-    (; p, dt) = integrator
-    (; graph, user_demand, basin, allocation) = p
-    (; flow, flow_dict, flow_prev, flow_integrated) = graph[]
-    (; vertical_flux, vertical_flux_prev, vertical_flux_integrated, vertical_flux_bmi) =
-        basin
-    flow = get_tmp(flow, 0)
-    vertical_flux = get_tmp(vertical_flux, 0)
-    if !isempty(flow_prev) && isnan(flow_prev[1])
-        # If flow_prev is not populated yet
-        copyto!(flow_prev, flow)
-    end
-
-    @. flow_integrated += 0.5 * (flow + flow_prev) * dt
-    @. vertical_flux_integrated += 0.5 * (vertical_flux + vertical_flux_prev) * dt
-    @. vertical_flux_bmi += 0.5 * (vertical_flux + vertical_flux_prev) * dt
-
-    # UserDemand realized flows for BMI
-    for (i, id) in enumerate(user_demand.node_id)
-        src_id = inflow_id(graph, id)
-        flow_idx = flow_dict[src_id, id]
-        user_demand.realized_bmi[i] += 0.5 * (flow[flow_idx] + flow_prev[flow_idx]) * dt
-    end
-
-    # Allocation source flows
-    for (edge, value) in allocation.mean_flows
-        if edge[1] == edge[2]
-            # Vertical fluxes
-            _, basin_idx = id_index(basin.node_id, edge[1])
-            value[] +=
-                0.5 *
-                (get_influx(basin, basin_idx) + get_influx(basin, basin_idx; prev = true)) *
-                dt
-        else
-            # Horizontal flows
-            value[] +=
-                0.5 *
-                (get_flow(graph, edge..., 0) + get_flow(graph, edge..., 0; prev = true)) *
-                dt
-        end
-    end
-
-    copyto!(flow_prev, flow)
-    copyto!(vertical_flux_prev, vertical_flux)
-    return nothing
-end
-
 "Compute the average flows over the last saveat interval and write
 them to SavedValues"
 function save_flow(u, t, integrator)
-    (; graph) = integrator.p
-    (; flow_integrated, flow_dict) = graph[]
+    (; uprev, p) = integrator
+    (; graph) = p
+    (; flow_dict) = graph[]
     (; node_id) = integrator.p.basin
 
     Δt = get_Δt(integrator)
-    flow_mean = copy(flow_integrated)
+    flow_mean = copy(u.flow_integrated)
     flow_mean ./= Δt
-    fill!(flow_integrated, 0.0)
+    fill!(u.flow_integrated, 0.0)
+    fill!(uprev.flow_integrated, 0.0)
 
     # Divide the flows over edges to Basin inflow and outflow, regardless of edge direction.
     inflow_mean = zeros(length(node_id))
@@ -177,22 +127,17 @@ end
 "Compute the average vertical fluxes over the last saveat interval and write
 them to SavedValues"
 function save_vertical_flux(u, t, integrator)
-    (; basin) = integrator.p
-    (; vertical_flux_integrated) = basin
-
     Δt = get_Δt(integrator)
-    vertical_flux_mean = copy(vertical_flux_integrated)
+    vertical_flux_mean = copy(forcings_integrated(u))
     vertical_flux_mean ./= Δt
-    fill!(vertical_flux_integrated, 0.0)
+    forcings_integrated(u) .= 0.0
 
     return vertical_flux_mean
 end
 
 function apply_discrete_control!(u, t, integrator)::Nothing
     (; p) = integrator
     (; discrete_control) = p
-    condition_idx = 0
-
     discrete_control_condition!(u, t, integrator)
 
     # For every compound variable see whether it changes a control state
@@ -432,7 +377,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)
 
@@ -453,17 +398,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) = allocation
 
     # Don't run the allocation algorithm if allocation is not active
     # (Specifically for running Ribasim via the BMI)
@@ -475,9 +417,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
+    u.flow_allocation_input /= Δt_allocation
 
     # If a main network is present, collect demands of subnetworks
     if has_main_network(allocation)
@@ -495,9 +435,9 @@ function update_allocation!(integrator)::Nothing
     end
 
     # Reset the mean source flows
-    for value in values(mean_flows)
-        value[] = 0.0
-    end
+    u.flow_allocation_input .= 0.0
+
+    return nothing
 end
 
 "Load updates from 'TabulatedRatingCurve / time' into the parameters"

core/src/graph.jl

@@ -91,20 +91,10 @@ function create_graph(db::DB, config::Config, chunk_sizes::Vector{Int})::MetaGra
     end
 
     flow = zeros(flow_counter)
-    flow_prev = fill(NaN, flow_counter)
-    flow_integrated = zeros(flow_counter)
     if config.solver.autodiff
         flow = DiffCache(flow, chunk_sizes)
     end
-    graph_data = (;
-        node_ids,
-        edges_source,
-        flow_dict,
-        flow,
-        flow_prev,
-        flow_integrated,
-        config.solver.saveat,
-    )
+    graph_data = (; node_ids, edges_source, flow_dict, flow, config.solver.saveat)
     graph = @set graph.graph_data = graph_data
 
     return graph
@@ -177,16 +167,9 @@ end
 """
 Get the flow over the given edge (val is needed for get_tmp from ForwardDiff.jl).
 """
-function get_flow(
-    graph::MetaGraph,
-    id_src::NodeID,
-    id_dst::NodeID,
-    val;
-    prev::Bool = false,
-)::Number
-    (; flow_dict, flow, flow_prev) = graph[]
-    flow_vector = prev ? flow_prev : flow
-    return get_tmp(flow_vector, val)[flow_dict[id_src, id_dst]]
+function get_flow(graph::MetaGraph, id_src::NodeID, id_dst::NodeID, val)::Number
+    (; flow_dict, flow) = graph[]
+    return get_tmp(flow, val)[flow_dict[id_src, id_dst]]
 end
 
 """