From 4e9f4793d0e3cd3ca2c90befe545742d59ad8fa2 Mon Sep 17 00:00:00 2001 From: Bart de Koning <74617371+SouthEndMusic@users.noreply.github.com> Date: Mon, 25 Nov 2024 12:52:35 +0100 Subject: [PATCH] Refactor of concentration code (#1929) Fixes https://github.com/Deltares/Ribasim/issues/1905 --------- Co-authored-by: Maarten Pronk --- core/src/Ribasim.jl | 3 +- core/src/callback.jl | 201 +++++++++++++------------------------- core/src/concentration.jl | 115 ++++++++++++++++++++++ core/src/graph.jl | 2 +- core/src/model.jl | 2 +- core/src/parameter.jl | 65 ++++++------ core/src/read.jl | 119 ++++++++++++---------- core/src/util.jl | 4 +- core/src/write.jl | 4 +- core/test/control_test.jl | 2 +- core/test/create_test.jl | 4 +- core/test/utils_test.jl | 22 ----- 12 files changed, 295 insertions(+), 248 deletions(-) create mode 100644 core/src/concentration.jl diff --git a/core/src/Ribasim.jl b/core/src/Ribasim.jl index b3656f4c1..f33f1045c 100644 --- a/core/src/Ribasim.jl +++ b/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 @@ -167,6 +167,7 @@ include("read.jl") include("write.jl") include("bmi.jl") include("callback.jl") +include("concentration.jl") include("main.jl") include("libribasim.jl") diff --git a/core/src/callback.jl b/core/src/callback.jl index ade71dc3f..8610564ce 100644 --- a/core/src/callback.jl +++ b/core/src/callback.jl @@ -34,6 +34,10 @@ function create_callbacks( FunctionCallingCallback(update_cumulative_flows!; func_start = false) push!(callbacks, cumulative_flows_cb) + # Update concentrations + concentrations_cb = FunctionCallingCallback(update_concentrations!; func_start = false) + push!(callbacks, concentrations_cb) + # Update Basin forcings tstops = get_tstops(basin.time.time, starttime) basin_cb = PresetTimeCallback(tstops, update_basin!; save_positions = (false, false)) @@ -113,30 +117,16 @@ Specifically, we first use all the inflows to update the mass of the Basins, rec the Basin concentration(s) and then remove the mass that is being lost to the outflows. """ function update_cumulative_flows!(u, t, integrator)::Nothing - (; p, uprev, tprev, dt) = integrator - (; - basin, - state_inflow_edge, - state_outflow_edge, - user_demand, - level_boundary, - flow_boundary, - allocation, - ) = p + (; p, tprev, dt) = integrator + (; basin, flow_boundary, allocation) = p (; vertical_flux) = basin # Update tprev p.tprev[] = t - # Reset cumulative flows, used to calculate the concentration - # of the basins after processing inflows only - basin.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 # Precipitation depends on fixed area for node_id in basin.node_id @@ -145,9 +135,6 @@ 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 end # Exact boundary flow over time step @@ -162,9 +149,6 @@ 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, :] .+= - flow_boundary.concentration[id.idx, :] .* volume - basin.cumulative_in[outflow_id.idx] += volume end end @@ -189,138 +173,83 @@ function update_cumulative_flows!(u, t, integrator)::Nothing end end end + return nothing +end - # Process mass updates for UserDemand separately - # as the inflow and outflow are decoupled in the states - for (inflow_edge, outflow_edge) in - zip(user_demand.inflow_edge, user_demand.outflow_edge) - from_node = inflow_edge.edge[1] - to_node = outflow_edge.edge[2] - userdemand_idx = outflow_edge.edge[1].idx - 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, :] .-= - 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, :] .+= - user_demand.concentration[userdemand_idx, :] .* flow - end - end - end +function update_concentrations!(u, t, integrator)::Nothing + (; uprev, p, tprev, dt) = integrator + (; basin, flow_boundary) = p + (; vertical_flux, concentration_data) = basin + (; evaporate_mass, cumulative_in, concentration_state, concentration, mass) = + concentration_data - # Process all mass inflows to basins - for (inflow_edge, outflow_edge) in zip(state_inflow_edge, state_outflow_edge) - from_node = inflow_edge.edge[1] - to_node = outflow_edge.edge[2] - 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 - if to_node.type == NodeType.Basin - basin.mass[from_node.idx, :] .-= - basin.concentration_state[to_node.idx, :] .* flow - elseif to_node.type == NodeType.LevelBoundary - basin.mass[from_node.idx, :] .-= - level_boundary.concentration[to_node.idx, :] .* flow - elseif to_node.type == NodeType.UserDemand - basin.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, - # and the unset flow edge defaults to Terminal #0 - nothing - else - @warn "Unsupported outflow from $(to_node.type) #$(to_node.value) to $(from_node.type) #$(from_node.value) with flow $flow" - end - end - end + # Reset cumulative flows, used to calculate the concentration + # of the basins after processing inflows only + cumulative_in .= 0.0 - 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 - if from_node.type == NodeType.Basin - basin.mass[to_node.idx, :] .+= - basin.concentration_state[from_node.idx, :] .* flow - elseif from_node.type == NodeType.LevelBoundary - basin.mass[to_node.idx, :] .+= - level_boundary.concentration[from_node.idx, :] .* flow - elseif from_node.type == NodeType.UserDemand - basin.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, - # and the unset flow edge defaults to Terminal #0 - nothing - else - @warn "Unsupported outflow from $(from_node.type) #$(from_node.value) to $(to_node.type) #$(to_node.value) with flow $flow" - end - end - end + mass .+= 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 + fixed_area = basin_areas(basin, node_id.idx)[end] + added_precipitation = fixed_area * vertical_flux.precipitation[node_id.idx] * dt + + mass[node_id.idx, :] .+= concentration[2, node_id.idx, :] .* added_precipitation + cumulative_in[node_id.idx] += added_precipitation end - # Update the Basin concentrations based on the added mass and flows - basin.concentration_state .= basin.mass ./ (basin.storage_prev .+ basin.cumulative_in) - - # Process all mass outflows from Basins - for (inflow_edge, outflow_edge) in zip(state_inflow_edge, state_outflow_edge) - from_node = inflow_edge.edge[1] - to_node = outflow_edge.edge[2] - 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 - 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 - end + # Exact boundary flow over time step + for (id, flow_rate, active, edge) in zip( + flow_boundary.node_id, + flow_boundary.flow_rate, + flow_boundary.active, + flow_boundary.outflow_edges, + ) + if active + outflow_id = edge[1].edge[2] + volume = integral(flow_rate, tprev, t) + mass[outflow_id.idx, :] .+= flow_boundary.concentration[id.idx, :] .* volume + cumulative_in[outflow_id.idx] += volume end end + mass_updates_user_demand!(integrator) + mass_inflows_basin!(integrator) + + # Update the Basin concentrations based on the added mass and flows + concentration_state .= mass ./ (basin.storage_prev .+ cumulative_in) + + mass_outflows_basin!(integrator) + # 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 evaporate_mass + mass .-= concentration_state .* (u.evaporation - uprev.evaporation) end - basin.mass .-= basin.concentration_state .* (u.infiltration - uprev.infiltration) + mass .-= 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)) < mass[I] < (eps(Float64)) + 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), mass) + R = CartesianIndices(mass) + locations = findall(<(0), 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)] + concentration_state .= 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)] - return nothing end @@ -427,7 +356,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 +611,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 +711,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 +733,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) diff --git a/core/src/concentration.jl b/core/src/concentration.jl new file mode 100644 index 000000000..12fb4dd2b --- /dev/null +++ b/core/src/concentration.jl @@ -0,0 +1,115 @@ +""" +Process mass updates for UserDemand separately +as the inflow and outflow are decoupled in the states +""" +function mass_updates_user_demand!(integrator::DEIntegrator)::Nothing + (; basin, user_demand) = integrator.p + (; concentration_state, mass) = basin.concentration_data + + for (inflow_edge, outflow_edge) in + zip(user_demand.inflow_edge, user_demand.outflow_edge) + from_node = inflow_edge.edge[1] + to_node = outflow_edge.edge[2] + userdemand_idx = outflow_edge.edge[1].idx + if from_node.type == NodeType.Basin + flow = flow_update_on_edge(integrator, inflow_edge.edge) + if flow < 0 + mass[from_node.idx, :] .-= concentration_state[to_node.idx, :] .* flow + 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 + mass[to_node.idx, :] .+= concentration_state[from_node.idx, :] .* flow + mass[to_node.idx, :] .+= + user_demand.concentration[userdemand_idx, :] .* flow + end + end + end + return nothing +end + +""" +Process all mass inflows to basins +""" +function mass_inflows_basin!(integrator::DEIntegrator)::Nothing + (; basin, state_inflow_edge, state_outflow_edge, level_boundary) = integrator.p + (; cumulative_in, concentration_state, mass) = basin.concentration_data + + for (inflow_edge, outflow_edge) in zip(state_inflow_edge, state_outflow_edge) + from_node = inflow_edge.edge[1] + to_node = outflow_edge.edge[2] + if from_node.type == NodeType.Basin + flow = flow_update_on_edge(integrator, inflow_edge.edge) + if flow < 0 + cumulative_in[from_node.idx] -= flow + if to_node.type == NodeType.Basin + mass[from_node.idx, :] .-= concentration_state[to_node.idx, :] .* flow + elseif to_node.type == NodeType.LevelBoundary + mass[from_node.idx, :] .-= + level_boundary.concentration[to_node.idx, :] .* flow + elseif to_node.type == NodeType.UserDemand + 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, + # and the unset flow edge defaults to Terminal #0 + nothing + else + @warn "Unsupported outflow from $(to_node.type) #$(to_node.value) to $(from_node.type) #$(from_node.value) with flow $flow" + end + end + end + + if to_node.type == NodeType.Basin + flow = flow_update_on_edge(integrator, outflow_edge.edge) + if flow > 0 + cumulative_in[to_node.idx] += flow + if from_node.type == NodeType.Basin + mass[to_node.idx, :] .+= concentration_state[from_node.idx, :] .* flow + elseif from_node.type == NodeType.LevelBoundary + mass[to_node.idx, :] .+= + level_boundary.concentration[from_node.idx, :] .* flow + elseif from_node.type == NodeType.UserDemand + 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, + # and the unset flow edge defaults to Terminal #0 + nothing + else + @warn "Unsupported outflow from $(from_node.type) #$(from_node.value) to $(to_node.type) #$(to_node.value) with flow $flow" + end + end + end + end + return nothing +end + +""" +Process all mass outflows from Basins +""" +function mass_outflows_basin!(integrator::DEIntegrator)::Nothing + (; state_inflow_edge, state_outflow_edge, basin) = integrator.p + (; mass, concentration_state) = basin.concentration_data + + for (inflow_edge, outflow_edge) in zip(state_inflow_edge, state_outflow_edge) + from_node = inflow_edge.edge[1] + to_node = outflow_edge.edge[2] + if from_node.type == NodeType.Basin + flow = flow_update_on_edge(integrator, inflow_edge.edge) + if flow > 0 + mass[from_node.idx, :] .-= 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 + mass[to_node.idx, :] .+= concentration_state[to_node.idx, :] .* flow + end + end + end + return nothing +end diff --git a/core/src/graph.jl b/core/src/graph.jl index d6c0669ba..d7e32d6a8 100644 --- a/core/src/graph.jl +++ b/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 diff --git a/core/src/model.jl b/core/src/model.jl index ccc9734bb..e0047ca0d 100644 --- a/core/src/model.jl +++ b/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) diff --git a/core/src/parameter.jl b/core/src/parameter.jl index 48a58e4a0..4e1c308d7 100644 --- a/core/src/parameter.jl +++ b/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 diff --git a/core/src/read.jl b/core/src/read.jl index 749554c3b..edcce108c 100644 --- a/core/src/read.jl +++ b/core/src/read.jl @@ -432,7 +432,7 @@ function LevelBoundary(db::DB, config::Config)::LevelBoundary concentration = zeros(length(node_ids), length(substances)) concentration[:, Substance.Continuity] .= 1.0 concentration[:, Substance.LevelBoundary] .= 1.0 - set_concentrations!(concentration, concentration_time, substances, Int32.(node_ids)) + set_concentrations!(concentration, concentration_time, substances, node_ids) if !valid error("Errors occurred when parsing LevelBoundary data.") @@ -475,7 +475,7 @@ function FlowBoundary(db::DB, config::Config, graph::MetaGraph)::FlowBoundary concentration = zeros(length(node_ids), length(substances)) concentration[:, Substance.Continuity] .= 1.0 concentration[:, Substance.FlowBoundary] .= 1.0 - set_concentrations!(concentration, concentration_time, substances, Int32.(node_ids)) + set_concentrations!(concentration, concentration_time, substances, node_ids) if !valid error("Errors occurred when parsing FlowBoundary data.") @@ -571,27 +571,17 @@ function Terminal(db::DB, config::Config)::Terminal return Terminal(NodeID.(NodeType.Terminal, node_id, eachindex(node_id))) end -function Basin(db::DB, config::Config, graph::MetaGraph)::Basin - node_id = get_ids(db, "Basin") +function ConcentrationData( + concentration_time, + node_id::Vector{NodeID}, + db::DB, + config::Config, +)::ConcentrationData n = length(node_id) - evaporate_mass = config.solver.evaporate_mass - precipitation = zeros(n) - potential_evaporation = zeros(n) - drainage = zeros(n) - infiltration = zeros(n) - table = (; precipitation, potential_evaporation, drainage, infiltration) - - area, level = create_storage_tables(db, config) - - # both static and time are optional, but we need fallback defaults - static = load_structvector(db, config, BasinStaticV1) - time = load_structvector(db, config, BasinTimeV1) - state = load_structvector(db, config, BasinStateV1) concentration_state_data = load_structvector(db, config, BasinConcentrationStateV1) - concentration_time = load_structvector(db, config, BasinConcentrationV1) - # TODO Move into a function + evaporate_mass = config.solver.evaporate_mass substances = get_substances(db, config) concentration_state = zeros(n, length(substances)) concentration_state[:, Substance.Continuity] .= 1.0 @@ -619,26 +609,6 @@ function Basin(db::DB, config::Config, graph::MetaGraph)::Basin concentration_column = :precipitation, ) - set_static_value!(table, node_id, static) - set_current_value!(table, node_id, time, config.starttime) - check_no_nans(table, "Basin") - - vertical_flux = - ComponentVector(; precipitation, potential_evaporation, drainage, infiltration) - - demand = zeros(length(node_id)) - - node_id = NodeID.(NodeType.Basin, node_id, eachindex(node_id)) - - is_valid = valid_profiles(node_id, level, area) - if !is_valid - error("Invalid Basin / profile table.") - end - - level_to_area = - LinearInterpolation.(area, level; extrapolate = true, cache_parameters = true) - storage_to_level = invert_integral.(level_to_area) - t_end = seconds_since(config.endtime, config.starttime) errors = false @@ -678,6 +648,61 @@ function Basin(db::DB, config::Config, graph::MetaGraph)::Basin error("Errors encountered when parsing Basin concentration data.") end + cumulative_in = zeros(n) + + return ConcentrationData(; + evaporate_mass, + concentration_state, + concentration, + mass, + concentration_external, + substances, + cumulative_in, + ) +end + +function Basin(db::DB, config::Config, graph::MetaGraph)::Basin + node_id = get_ids(db, "Basin") + n = length(node_id) + + # both static and time are optional, but we need fallback defaults + static = load_structvector(db, config, BasinStaticV1) + time = load_structvector(db, config, BasinTimeV1) + state = load_structvector(db, config, BasinStateV1) + + # Forcing + precipitation = zeros(n) + potential_evaporation = zeros(n) + drainage = zeros(n) + infiltration = zeros(n) + table = (; precipitation, potential_evaporation, drainage, infiltration) + + set_static_value!(table, node_id, static) + set_current_value!(table, node_id, time, config.starttime) + check_no_nans(table, "Basin") + + vertical_flux = ComponentVector(; table...) + + # Node IDs + node_id = NodeID.(NodeType.Basin, node_id, eachindex(node_id)) + + # Profiles + area, level = create_storage_tables(db, config) + + is_valid = valid_profiles(node_id, level, area) + if !is_valid + error("Invalid Basin / profile table.") + end + + level_to_area = + LinearInterpolation.(area, level; extrapolate = true, cache_parameters = true) + storage_to_level = invert_integral.(level_to_area) + + # Concentration data + concentration_time = load_structvector(db, config, BasinConcentrationV1) + concentration_data = ConcentrationData(concentration_time, node_id, db, config) + + # Initialize Basin basin = Basin(; node_id, inflow_ids = [collect(inflow_ids(graph, id)) for id in node_id], @@ -685,22 +710,16 @@ function Basin(db::DB, config::Config, graph::MetaGraph)::Basin vertical_flux, storage_to_level, level_to_area, - demand, time, + concentration_data, concentration_time, - evaporate_mass, - concentration_state, - concentration, - mass, - concentration_external, - substances, ) storage0 = get_storages_from_levels(basin, state.level) @assert length(storage0) == n "Basin / state length differs from number of Basins" basin.storage0 .= storage0 basin.storage_prev .= storage0 - basin.mass .*= storage0 # was initialized by concentration_state, resulting in mass + basin.concentration_data.mass .*= storage0 # was initialized by concentration_state, resulting in mass return basin end @@ -1128,7 +1147,7 @@ function UserDemand(db::DB, config::Config, graph::MetaGraph)::UserDemand concentration = zeros(length(node_ids), length(substances)) # Continuity concentration is zero, as the return flow (from a Basin) already includes it concentration[:, Substance.UserDemand] .= 1.0 - set_concentrations!(concentration, concentration_time, substances, ids) + set_concentrations!(concentration, concentration_time, substances, node_ids) if errors || !valid_demand(node_ids, demand_itp, priorities) error("Errors occurred when parsing UserDemand data.") @@ -1521,7 +1540,7 @@ function set_concentrations!( concentration, concentration_data, substances, - node_ids; + node_ids::Vector{NodeID}; concentration_column = :concentration, ) for substance in unique(concentration_data.substance) @@ -1531,7 +1550,7 @@ function set_concentrations!( first_row = first(group) value = getproperty(first_row, concentration_column) ismissing(value) && continue - node_idx = findfirst(==(first_row.node_id), node_ids) + node_idx = findfirst(node_id -> node_id.value == first_row.node_id, node_ids) concentration[node_idx, sub_idx] = value end end diff --git a/core/src/util.jl b/core/src/util.jl index 0ccfddb28..c05057911 100644 --- a/core/src/util.jl +++ b/core/src/util.jl @@ -1059,8 +1059,8 @@ function set_state_flow_edges(p::Parameters, u0::ComponentVector)::Parameters state_inflow_edge = ComponentVector(NamedTuple(zip(components, state_inflow_edges))) state_outflow_edge = ComponentVector(NamedTuple(zip(components, state_outflow_edges))) - p = @set p.state_inflow_edge = state_inflow_edge - p = @set p.state_outflow_edge = state_outflow_edge + @reset p.state_inflow_edge = state_inflow_edge + @reset p.state_outflow_edge = state_outflow_edge return p end diff --git a/core/src/write.jl b/core/src/write.jl index bb3ad0472..b81e5d01d 100644 --- a/core/src/write.jl +++ b/core/src/write.jl @@ -275,9 +275,9 @@ function concentration_table( ntsteps = length(data.time) - 1 nbasin = length(data.node_id) - nsubstance = length(basin.substances) + nsubstance = length(basin.concentration_data.substances) - substances = String.(basin.substances) + substances = String.(basin.concentration_data.substances) concentration = FlatVector(saved.flow.saveval, :concentration) time = repeat(data.time[begin:(end - 1)]; inner = nbasin * nsubstance) diff --git a/core/test/control_test.jl b/core/test/control_test.jl index a9efea9cb..737a8da82 100644 --- a/core/test/control_test.jl +++ b/core/test/control_test.jl @@ -297,7 +297,7 @@ end flow_edge_0 = filter(:edge_id => id -> id == 0, flow_data) t = Ribasim.seconds_since.(flow_edge_0.time, model.config.starttime) itp = - model.integrator.p.basin.concentration_external[1]["concentration_external.kryptonite"] + model.integrator.p.basin.concentration_data.concentration_external[1]["concentration_external.kryptonite"] concentration = itp.(t) threshold = 0.5 above_threshold = concentration .> threshold diff --git a/core/test/create_test.jl b/core/test/create_test.jl index 90b4d1cc5..01865b3f2 100644 --- a/core/test/create_test.jl +++ b/core/test/create_test.jl @@ -3,7 +3,7 @@ using Graphs using Logging using Ribasim: NodeID - using Accessors: @set + using Accessors: @set, @reset graph = MetaGraph( DiGraph(); @@ -25,7 +25,7 @@ push!(node_ids[-1], NodeID(:Basin, 2, 1)) graph_data = (; node_ids,) - graph = @set graph.graph_data = graph_data + @reset graph.graph_data = graph_data logger = TestLogger() with_logger(logger) do diff --git a/core/test/utils_test.jl b/core/test/utils_test.jl index 74e082f08..0b8e2c3f9 100644 --- a/core/test/utils_test.jl +++ b/core/test/utils_test.jl @@ -20,22 +20,11 @@ end level = [[0.0, 1.0], [4.0, 5.0]] level_to_area = LinearInterpolation.(area, level) storage_to_level = invert_integral.(level_to_area) - demand = zeros(2) - - substances = OrderedSet([:test]) - concentration_state = zeros(2, 1) - concentration = zeros(2, 2, 1) - mass = zeros(2, 1) basin = Ribasim.Basin(; node_id = NodeID.(:Basin, [5, 7], [1, 2]), storage_to_level, level_to_area, - demand, - concentration_state, - concentration, - mass, - substances, time = StructVector{Ribasim.BasinTimeV1}(undef, 0), concentration_time = StructVector{Ribasim.BasinConcentrationV1}(undef, 0), ) @@ -84,24 +73,13 @@ end ] level_to_area = LinearInterpolation(area, level; extrapolate = true) storage_to_level = invert_integral(level_to_area) - demand = zeros(1) - - substances = OrderedSet([:test]) - concentration_state = zeros(1, 1) - concentration = zeros(2, 1, 1) - mass = zeros(1, 1) basin = Ribasim.Basin(; node_id = NodeID.(:Basin, [1], 1), storage_to_level = [storage_to_level], level_to_area = [level_to_area], - demand, time = StructVector{Ribasim.BasinTimeV1}(undef, 0), concentration_time = StructVector{Ribasim.BasinConcentrationV1}(undef, 0), - concentration_state, - concentration, - mass, - substances, ) logger = TestLogger()