Refactor of concentration code

core/src/Ribasim.jl

@@ -132,7 +132,7 @@ using MetaGraphsNext:
 using EnumX: EnumX, @enumx
 
 # Easily change an immutable field of an object.
-using Accessors: @set
+using Accessors: @set, @reset
 
 # Iteration utilities, used to partition and group tables.
 import IterTools

core/src/callback.jl

@@ -130,13 +130,14 @@ function update_cumulative_flows!(u, t, integrator)::Nothing
 
     # Reset cumulative flows, used to calculate the concentration
     # of the basins after processing inflows only
-    basin.cumulative_in .= 0.0
+    basin.concentration_data.cumulative_in .= 0.0
 
     # Update cumulative forcings which are integrated exactly
     @. basin.cumulative_drainage += vertical_flux.drainage * dt
     @. basin.cumulative_drainage_saveat += vertical_flux.drainage * dt
-    basin.mass .+= basin.concentration[1, :, :] .* vertical_flux.drainage * dt
-    basin.cumulative_in .= vertical_flux.drainage * dt
+    basin.concentration_data.mass .+=
+        basin.concentration_data.concentration[1, :, :] .* vertical_flux.drainage * dt
+    basin.concentration_data.cumulative_in .= vertical_flux.drainage * dt
 
     # Precipitation depends on fixed area
     for node_id in basin.node_id
@@ -145,9 +146,9 @@ function update_cumulative_flows!(u, t, integrator)::Nothing
 
         basin.cumulative_precipitation[node_id.idx] += added_precipitation
         basin.cumulative_precipitation_saveat[node_id.idx] += added_precipitation
-        basin.mass[node_id.idx, :] .+=
-            basin.concentration[2, node_id.idx, :] .* added_precipitation
-        basin.cumulative_in[node_id.idx] += added_precipitation
+        basin.concentration_data.mass[node_id.idx, :] .+=
+            basin.concentration_data.concentration[2, node_id.idx, :] .* added_precipitation
+        basin.concentration_data.cumulative_in[node_id.idx] += added_precipitation
     end
 
     # Exact boundary flow over time step
@@ -162,9 +163,9 @@ function update_cumulative_flows!(u, t, integrator)::Nothing
             volume = integral(flow_rate, tprev, t)
             flow_boundary.cumulative_flow[id.idx] += volume
             flow_boundary.cumulative_flow_saveat[id.idx] += volume
-            basin.mass[outflow_id.idx, :] .+=
+            basin.concentration_data.mass[outflow_id.idx, :] .+=
                 flow_boundary.concentration[id.idx, :] .* volume
-            basin.cumulative_in[outflow_id.idx] += volume
+            basin.concentration_data.cumulative_in[outflow_id.idx] += volume
         end
     end
 
@@ -200,18 +201,18 @@ function update_cumulative_flows!(u, t, integrator)::Nothing
         if from_node.type == NodeType.Basin
             flow = flow_update_on_edge(integrator, inflow_edge.edge)
             if flow < 0
-                basin.mass[from_node.idx, :] .-=
-                    basin.concentration_state[to_node.idx, :] .* flow
-                basin.mass[from_node.idx, :] .-=
+                basin.concentration_data.mass[from_node.idx, :] .-=
+                    basin.concentration_data.concentration_state[to_node.idx, :] .* flow
+                basin.concentration_data.mass[from_node.idx, :] .-=
                     user_demand.concentration[userdemand_idx, :] .* flow
             end
         end
         if to_node.type == NodeType.Basin
             flow = flow_update_on_edge(integrator, outflow_edge.edge)
             if flow > 0
-                basin.mass[to_node.idx, :] .+=
-                    basin.concentration_state[from_node.idx, :] .* flow
-                basin.mass[to_node.idx, :] .+=
+                basin.concentration_data.mass[to_node.idx, :] .+=
+                    basin.concentration_data.concentration_state[from_node.idx, :] .* flow
+                basin.concentration_data.mass[to_node.idx, :] .+=
                     user_demand.concentration[userdemand_idx, :] .* flow
             end
         end
@@ -224,15 +225,15 @@ function update_cumulative_flows!(u, t, integrator)::Nothing
         if from_node.type == NodeType.Basin
             flow = flow_update_on_edge(integrator, inflow_edge.edge)
             if flow < 0
-                basin.cumulative_in[from_node.idx] -= flow
+                basin.concentration_data.cumulative_in[from_node.idx] -= flow
                 if to_node.type == NodeType.Basin
-                    basin.mass[from_node.idx, :] .-=
-                        basin.concentration_state[to_node.idx, :] .* flow
+                    basin.concentration_data.mass[from_node.idx, :] .-=
+                        basin.concentration_data.concentration_state[to_node.idx, :] .* flow
                 elseif to_node.type == NodeType.LevelBoundary
-                    basin.mass[from_node.idx, :] .-=
+                    basin.concentration_data.mass[from_node.idx, :] .-=
                         level_boundary.concentration[to_node.idx, :] .* flow
                 elseif to_node.type == NodeType.UserDemand
-                    basin.mass[from_node.idx, :] .-=
+                    basin.concentration_data.mass[from_node.idx, :] .-=
                         user_demand.concentration[to_node.idx, :] .* flow
                 elseif to_node.type == NodeType.Terminal && to_node.value == 0
                     # UserDemand inflow is discoupled from its outflow,
@@ -247,15 +248,16 @@ function update_cumulative_flows!(u, t, integrator)::Nothing
         if to_node.type == NodeType.Basin
             flow = flow_update_on_edge(integrator, outflow_edge.edge)
             if flow > 0
-                basin.cumulative_in[to_node.idx] += flow
+                basin.concentration_data.cumulative_in[to_node.idx] += flow
                 if from_node.type == NodeType.Basin
-                    basin.mass[to_node.idx, :] .+=
-                        basin.concentration_state[from_node.idx, :] .* flow
+                    basin.concentration_data.mass[to_node.idx, :] .+=
+                        basin.concentration_data.concentration_state[from_node.idx, :] .*
+                        flow
                 elseif from_node.type == NodeType.LevelBoundary
-                    basin.mass[to_node.idx, :] .+=
+                    basin.concentration_data.mass[to_node.idx, :] .+=
                         level_boundary.concentration[from_node.idx, :] .* flow
                 elseif from_node.type == NodeType.UserDemand
-                    basin.mass[to_node.idx, :] .+=
+                    basin.concentration_data.mass[to_node.idx, :] .+=
                         user_demand.concentration[from_node.idx, :] .* flow
                 elseif from_node.type == NodeType.Terminal && from_node.value == 0
                     # UserDemand outflow is discoupled from its inflow,
@@ -269,7 +271,9 @@ function update_cumulative_flows!(u, t, integrator)::Nothing
     end
 
     # Update the Basin concentrations based on the added mass and flows
-    basin.concentration_state .= basin.mass ./ (basin.storage_prev .+ basin.cumulative_in)
+    basin.concentration_data.concentration_state .=
+        basin.concentration_data.mass ./
+        (basin.storage_prev .+ basin.concentration_data.cumulative_in)
 
     # Process all mass outflows from Basins
     for (inflow_edge, outflow_edge) in zip(state_inflow_edge, state_outflow_edge)
@@ -278,46 +282,50 @@ function update_cumulative_flows!(u, t, integrator)::Nothing
         if from_node.type == NodeType.Basin
             flow = flow_update_on_edge(integrator, inflow_edge.edge)
             if flow > 0
-                basin.mass[from_node.idx, :] .-=
-                    basin.concentration_state[from_node.idx, :] .* flow
+                basin.concentration_data.mass[from_node.idx, :] .-=
+                    basin.concentration_data.concentration_state[from_node.idx, :] .* flow
             end
         end
         if to_node.type == NodeType.Basin
             flow = flow_update_on_edge(integrator, outflow_edge.edge)
             if flow < 0
-                basin.mass[to_node.idx, :] .+=
-                    basin.concentration_state[to_node.idx, :] .* flow
+                basin.concentration_data.mass[to_node.idx, :] .+=
+                    basin.concentration_data.concentration_state[to_node.idx, :] .* flow
             end
         end
     end
 
     # Evaporate mass to keep the mass balance, if enabled in model config
-    if basin.evaporate_mass
-        basin.mass .-= basin.concentration_state .* (u.evaporation - uprev.evaporation)
+    if basin.concentration_data.evaporate_mass
+        basin.concentration_data.mass .-=
+            basin.concentration_data.concentration_state .*
+            (u.evaporation - uprev.evaporation)
     end
-    basin.mass .-= basin.concentration_state .* (u.infiltration - uprev.infiltration)
+    basin.concentration_data.mass .-=
+        basin.concentration_data.concentration_state .*
+        (u.infiltration - uprev.infiltration)
 
     # Take care of infinitely small masses, possibly becoming negative due to truncation.
-    for I in eachindex(basin.mass)
-        if (-eps(Float64)) < basin.mass[I] < (eps(Float64))
-            basin.mass[I] = 0.0
+    for I in eachindex(basin.concentration_data.mass)
+        if (-eps(Float64)) < basin.concentration_data.mass[I] < (eps(Float64))
+            basin.concentration_data.mass[I] = 0.0
         end
     end
 
     # Check for negative masses
-    if any(<(0), basin.mass)
-        R = CartesianIndices(basin.mass)
-        locations = findall(<(0), basin.mass)
+    if any(<(0), basin.concentration_data.mass)
+        R = CartesianIndices(basin.concentration_data.mass)
+        locations = findall(<(0), basin.concentration_data.mass)
         for I in locations
             basin_idx, substance_idx = Tuple(R[I])
-            @error "$(basin.node_id[basin_idx]) has negative mass $(basin.mass[I]) for substance $(basin.substances[substance_idx])"
+            @error "$(basin.node_id[basin_idx]) has negative mass $(basin.concentration_data.mass[I]) for substance $(basin.concentration_data.substances[substance_idx])"
         end
         error("Negative mass(es) detected")
     end
 
     # Update the Basin concentrations again based on the removed mass
-    basin.concentration_state .=
-        basin.mass ./ basin.current_properties.current_storage[parent(u)]
+    basin.concentration_data.concentration_state .=
+        basin.concentration_data.mass ./ basin.current_properties.current_storage[parent(u)]
     basin.storage_prev .= basin.current_properties.current_storage[parent(u)]
     basin.level_prev .= basin.current_properties.current_level[parent(u)]
 
@@ -427,7 +435,7 @@ function save_flow(u, t, integrator)
     @. basin.cumulative_precipitation_saveat = 0.0
     @. basin.cumulative_drainage_saveat = 0.0
 
-    concentration = copy(basin.concentration_state)
+    concentration = copy(basin.concentration_data.concentration_state)
     saved_flow = SavedFlow(;
         flow = flow_mean,
         inflow = inflow_mean,
@@ -682,11 +690,12 @@ function get_value(subvariable::NamedTuple, p::Parameters, du::AbstractVector, t
         end
 
     elseif startswith(variable, "concentration_external.")
-        value = basin.concentration_external[listen_node_id.idx][variable](t)
+        value =
+            basin.concentration_data.concentration_external[listen_node_id.idx][variable](t)
     elseif startswith(variable, "concentration.")
         substance = Symbol(last(split(variable, ".")))
-        var_idx = findfirst(==(substance), basin.substances)
-        value = basin.concentration_state[listen_node_id.idx, var_idx]
+        var_idx = findfirst(==(substance), basin.concentration_data.substances)
+        value = basin.concentration_data.concentration_state[listen_node_id.idx, var_idx]
     else
         error("Unsupported condition variable $variable.")
     end
@@ -781,7 +790,8 @@ end
 function update_basin_conc!(integrator)::Nothing
     (; p) = integrator
     (; basin) = p
-    (; node_id, concentration, concentration_time, substances) = basin
+    (; node_id, concentration_data, concentration_time) = basin
+    (; concentration, substances) = concentration_data
     t = datetime_since(integrator.t, integrator.p.starttime)
 
     rows = searchsorted(concentration_time.time, t)
@@ -802,7 +812,7 @@ function update_conc!(integrator, parameter, nodetype)::Nothing
     node = getproperty(p, parameter)
     (; basin) = p
     (; node_id, concentration, concentration_time) = node
-    (; substances) = basin
+    (; substances) = basin.concentration_data
     t = datetime_since(integrator.t, integrator.p.starttime)
 
     rows = searchsorted(concentration_time.time, t)

core/src/graph.jl

@@ -107,7 +107,7 @@ function create_graph(db::DB, config::Config)::MetaGraph
     end
 
     graph_data = (; node_ids, edges_source, flow_edges, config.solver.saveat)
-    graph = @set graph.graph_data = graph_data
+    @reset graph.graph_data = graph_data
 
     return graph
 end

core/src/model.jl

@@ -105,7 +105,7 @@ function Model(config::Config)::Model
 
     parameters = set_state_flow_edges(parameters, u0)
     parameters = build_flow_to_storage(parameters, u0)
-    parameters = @set parameters.u_prev_saveat = zero(u0)
+    @reset parameters.u_prev_saveat = zero(u0)
 
     # The Solver algorithm
     alg = algorithm(config.solver; u0)

core/src/parameter.jl

@@ -324,6 +324,24 @@ struct CurrentBasinProperties
     end
 end
 
+@kwdef struct ConcentrationData
+    # Config setting to enable/disable evaporation of mass
+    evaporate_mass::Bool = true
+    # Cumulative inflow for each Basin at a given time
+    cumulative_in::Vector{Float64}
+    # matrix with concentrations for each Basin and substance
+    concentration_state::Matrix{Float64}  # Basin, substance
+    # matrix with boundary concentrations for each boundary, Basin and substance
+    concentration::Array{Float64, 3}
+    # matrix with mass for each Basin and substance
+    mass::Matrix{Float64}
+    # substances in use by the model (ordered like their axis in the concentration matrices)
+    substances::OrderedSet{Symbol}
+    # Data source for external concentrations (used in control)
+    concentration_external::Vector{Dict{String, ScalarInterpolation}} =
+        Dict{String, ScalarInterpolation}[]
+end
+
 """
 Requirements:
 
@@ -340,7 +358,7 @@ else
     T = Vector{Float64}
 end
 """
-@kwdef struct Basin{C, D, V} <: AbstractParameterNode
+@kwdef struct Basin{V, C, CD, D} <: AbstractParameterNode
     node_id::Vector{NodeID}
     inflow_ids::Vector{Vector{NodeID}} = [NodeID[]]
     outflow_ids::Vector{Vector{NodeID}} = [NodeID[]]
@@ -369,32 +387,17 @@ end
     }
     level_to_area::Vector{ScalarInterpolation}
     # Demands for allocation if applicable
-    demand::Vector{Float64}
+    demand::Vector{Float64} = zeros(length(node_id))
     # Data source for parameter updates
     time::StructVector{BasinTimeV1, C, Int}
-    # Data source for concentration updates
-    concentration_time::StructVector{BasinConcentrationV1, D, Int}
-
+    # Storage for each Basin at the previous time step
+    storage_prev::Vector{Float64} = zeros(length(node_id))
     # Level for each Basin at the previous time step
     level_prev::Vector{Float64} = zeros(length(node_id))
     # Concentrations
-    # Config setting to enable/disable evaporation of mass
-    evaporate_mass::Bool = true
-    # Cumulative inflow for each Basin at a given time
-    cumulative_in::Vector{Float64} = zeros(length(node_id))
-    # Storage for each Basin at the previous time step
-    storage_prev::Vector{Float64} = zeros(length(node_id))
-    # matrix with concentrations for each Basin and substance
-    concentration_state::Matrix{Float64}  # Basin, substance
-    # matrix with boundary concentrations for each boundary, Basin and substance
-    concentration::Array{Float64, 3}
-    # matrix with mass for each Basin and substance
-    mass::Matrix{Float64}
-    # substances in use by the model (ordered like their axis in the concentration matrices)
-    substances::OrderedSet{Symbol}
-    # Data source for external concentrations (used in control)
-    concentration_external::Vector{Dict{String, ScalarInterpolation}} =
-        Dict{String, ScalarInterpolation}[]
+    concentration_data::CD = nothing
+    # Data source for concentration updates
+    concentration_time::StructVector{BasinConcentrationV1, D, Int}
 end
 
 """
@@ -880,29 +883,29 @@ const ModelGraph = MetaGraph{
     Float64,
 }
 
-@kwdef struct Parameters{C1, C2, C3, C4, C5, C6, C7, C8, C9, V}
+@kwdef struct Parameters{C1, C2, C3, C4, C5, C6, C7, C8, C9, C10, C11}
     starttime::DateTime
     graph::ModelGraph
     allocation::Allocation
-    basin::Basin{C1, C2, V}
+    basin::Basin{C1, C2, C3, C4}
     linear_resistance::LinearResistance
     manning_resistance::ManningResistance
-    tabulated_rating_curve::TabulatedRatingCurve{C3}
-    level_boundary::LevelBoundary{C4}
-    flow_boundary::FlowBoundary{C5}
+    tabulated_rating_curve::TabulatedRatingCurve{C5}
+    level_boundary::LevelBoundary{C6}
+    flow_boundary::FlowBoundary{C7}
     pump::Pump
     outlet::Outlet
     terminal::Terminal
     discrete_control::DiscreteControl
     continuous_control::ContinuousControl
     pid_control::PidControl
-    user_demand::UserDemand{C6}
+    user_demand::UserDemand{C8}
     level_demand::LevelDemand
     flow_demand::FlowDemand
     subgrid::Subgrid
     # Per state the in- and outflow edges associated with that state (if they exist)
-    state_inflow_edge::C7 = ComponentVector()
-    state_outflow_edge::C8 = ComponentVector()
+    state_inflow_edge::C9 = ComponentVector()
+    state_outflow_edge::C10 = ComponentVector()
     all_nodes_active::Base.RefValue{Bool} = Ref(false)
     tprev::Base.RefValue{Float64} = Ref(0.0)
     # Sparse matrix for combining flows into storages
@@ -911,7 +914,7 @@ const ModelGraph = MetaGraph{
     water_balance_abstol::Float64
     water_balance_reltol::Float64
     # State at previous saveat
-    u_prev_saveat::C9 = ComponentVector()
+    u_prev_saveat::C11 = ComponentVector()
 end
 
 # To opt-out of type checking for ForwardDiff