From 6c0c7ee8dd07fd1c1e3e01d137be4de028173696 Mon Sep 17 00:00:00 2001 From: Maarten Pronk Date: Tue, 31 Mar 2026 21:33:21 +0200 Subject: [PATCH 1/2] Add subtimestepping for the internal tracer calculation. --- core/src/callback.jl | 234 ++++++++++++++++++++++---------------- core/src/concentration.jl | 60 +++++----- core/src/config.jl | 8 +- core/src/parameter.jl | 9 ++ core/src/read.jl | 8 +- core/src/write.jl | 3 + 6 files changed, 198 insertions(+), 124 deletions(-) diff --git a/core/src/callback.jl b/core/src/callback.jl index a20e863a00..1ac4e04e09 100644 --- a/core/src/callback.jl +++ b/core/src/callback.jl @@ -189,8 +189,12 @@ function update_concentrations!(u, t, integrator)::Nothing (; vertical_flux, concentration_data) = basin (; evaporate_mass, + substep_ratio, + substep_depth, cumulative_in, concentration_state, + nsubsteps, + stepsize, concentration_itp_drainage, concentration_itp_precipitation, concentration_itp_surface_runoff, @@ -200,125 +204,157 @@ function update_concentrations!(u, t, integrator)::Nothing !do_concentration && return nothing - # Reset cumulative flows, used to calculate the concentration - cumulative_in .= vertical_flux.drainage * dt - cumulative_in .+= vertical_flux.surface_runoff * dt + # Determine number of substeps needed based on a fraction of the residence time + safe = x -> isinf(x) ? 0 : x + @. nsubsteps = 2^clamp(floor(Int, safe(log2(dt / (concentration_state[:, Substance.ResidenceTime] * substep_ratio)))), 0, substep_depth) + max_substeps = maximum(nsubsteps) + @. stepsize = max_substeps รท nsubsteps - # Basin forcings - for node_id in basin.node_id - mass_node = mass[node_id.idx] + dt_sub = dt / max_substeps - add_substance_mass!( - mass_node, - concentration_itp_drainage[node_id.idx], - vertical_flux.drainage[node_id.idx] * dt, - t, - ) + for substep in 1:max_substeps - # Precipitation depends on fixed area - fixed_area = basin_areas(basin, node_id.idx)[end] - added_precipitation = fixed_area * vertical_flux.precipitation[node_id.idx] * dt - add_substance_mass!( - mass_node, - concentration_itp_precipitation[node_id.idx], - added_precipitation, - t, - ) - cumulative_in[node_id.idx] += added_precipitation + # Basin forcings + for node_id in basin.node_id + # Check if we need to compute the current substep for this node + (substep % stepsize[node_id.idx]) != 0 && continue - add_substance_mass!( - mass_node, - concentration_itp_surface_runoff[node_id.idx], - vertical_flux.surface_runoff[node_id.idx] * dt, - t, - ) + # Initialize cumulative_in for this basin's processing window + cumulative_in[node_id.idx] = (vertical_flux.drainage[node_id.idx] + vertical_flux.surface_runoff[node_id.idx]) * dt_sub * stepsize[node_id.idx] + mass_node = mass[node_id.idx] - add_substance_mass!( - mass_node, - loads_itp[node_id.idx], - dt, # loads are per second, not volume, so the flow is just the time step - t, - ) - end + add_substance_mass!( + mass_node, + concentration_itp_drainage[node_id.idx], + vertical_flux.drainage[node_id.idx] * dt_sub * stepsize[node_id.idx], + tprev + dt_sub * substep, + ) - # Exact boundary flow over time step - for (id, flow_rate, outflow_link) in zip( - flow_boundary.node_id, - flow_boundary.flow_rate, - flow_boundary.outflow_link, - ) - outflow_id = outflow_link.link[2] - added_boundary_flow = integral(flow_rate, tprev, t) - add_substance_mass!( - mass[outflow_id.idx], - flow_boundary.concentration_itp[id.idx], - added_boundary_flow, - t, - ) - cumulative_in[outflow_id.idx] += added_boundary_flow - end + # Precipitation depends on fixed area + fixed_area = basin_areas(basin, node_id.idx)[end] + added_precipitation = fixed_area * vertical_flux.precipitation[node_id.idx] * dt_sub * stepsize[node_id.idx] + add_substance_mass!( + mass_node, + concentration_itp_precipitation[node_id.idx], + added_precipitation, + tprev + dt_sub * substep, + ) + cumulative_in[node_id.idx] += added_precipitation - mass_inflows_from_user_demand!(integrator) - mass_inflows_basin!(integrator) + add_substance_mass!( + mass_node, + concentration_itp_surface_runoff[node_id.idx], + vertical_flux.surface_runoff[node_id.idx] * dt_sub * stepsize[node_id.idx], + tprev + dt_sub * substep, + ) - # Update the Basin concentrations based on the added mass and flows - for node_id in basin.node_id - storage_only_in = basin.storage_prev[node_id.idx] + cumulative_in[node_id.idx] + add_substance_mass!( + mass_node, + loads_itp[node_id.idx], + dt_sub * stepsize[node_id.idx], # loads are per second, not volume, so the flow is just the time step + tprev + dt_sub * substep, + ) - # The residence time tracer gets older - mass[node_id.idx][Substance.ResidenceTime] += dt * basin.storage_prev[node_id.idx] - if iszero(storage_only_in) - concentration_state[node_id.idx, :] .= 0 - else - concentration_state[node_id.idx, :] .= mass[node_id.idx] ./ storage_only_in end - end - mass_outflows_basin!(integrator) + # Exact boundary flow over time step + for (id, flow_rate, outflow_link) in zip( + flow_boundary.node_id, + flow_boundary.flow_rate, + flow_boundary.outflow_link, + ) + outflow_id = outflow_link.link[2] + (substep % stepsize[outflow_id.idx]) != 0 && continue + + added_boundary_flow = integral(flow_rate, tprev + (substep - stepsize[outflow_id.idx]) * dt_sub, tprev + substep * dt_sub) + add_substance_mass!( + mass[outflow_id.idx], + flow_boundary.concentration_itp[id.idx], + added_boundary_flow, + tprev + dt_sub * substep, + ) + cumulative_in[outflow_id.idx] += added_boundary_flow + end - errors = false + mass_inflows_from_user_demand!(integrator, substep, max_substeps) + mass_inflows_basin!(integrator, substep, max_substeps) - for node_id in basin.node_id - mass_node = mass[node_id.idx] + # Update the Basin concentrations based on the added mass and flows + for node_id in basin.node_id + (substep % stepsize[node_id.idx]) != 0 && continue - # Evaporate mass to keep the mass balance, if enabled in model config - if evaporate_mass - evaporated_volume = u.evaporation[node_id.idx] - uprev.evaporation[node_id.idx] - mass_node .-= concentration_state[node_id.idx, :] .* evaporated_volume - end + # Storage at the start of this processing window + inflows during the window + s_start = basin.storage_prev[node_id.idx] + (current_storage[node_id.idx] - basin.storage_prev[node_id.idx]) * (substep - stepsize[node_id.idx]) / max_substeps + storage_only_in = s_start + cumulative_in[node_id.idx] - infiltrated_volume = u.infiltration[node_id.idx] - uprev.infiltration[node_id.idx] - mass_node .-= concentration_state[node_id.idx, :] .* infiltrated_volume + # The residence time tracer gets older + mass[node_id.idx][Substance.ResidenceTime] += dt_sub * stepsize[node_id.idx] * basin.storage_prev[node_id.idx] + if iszero(storage_only_in) + concentration_state[node_id.idx, :] .= 0 + else + concentration_state[node_id.idx, :] .= mass[node_id.idx] ./ storage_only_in + end - # Take care of infinitely small masses, possibly becoming negative due to truncation. - for I in eachindex(mass_node) - if (-eps(Float64)) < mass_node[I] < (eps(Float64)) - mass_node[I] = 0.0 + # Debug: check continuity tracer after dilution step + c_cont = concentration_state[node_id.idx, Substance.Continuity] + if !iszero(storage_only_in) && abs(c_cont - 1.0) > 0.01 + @warn "Continuity drift after dilution" node_id substep c_cont storage_only_in cumulative_in = cumulative_in[node_id.idx] s_start mass_cont = mass[node_id.idx][Substance.Continuity] end end - # Check for negative masses - if any(<(0), mass_node) - errors = true - for substance_idx in findall(<(0), mass_node) - substance_name = basin.concentration_data.substances[substance_idx] - substance_mass = mass_node[substance_idx] - @error "$node_id has negative mass $substance_mass for substance $substance_name" + mass_outflows_basin!(integrator, substep, max_substeps) + + errors = false + + for node_id in basin.node_id + (substep % stepsize[node_id.idx]) != 0 && continue + + mass_node = mass[node_id.idx] + + # Evaporate mass to keep the mass balance, if enabled in model config + if evaporate_mass + evaporated_volume = (u.evaporation[node_id.idx] - uprev.evaporation[node_id.idx]) / nsubsteps[node_id.idx] + mass_node .-= concentration_state[node_id.idx, :] .* evaporated_volume end - end - # Update the Basin concentrations again based on the removed mass - s = current_storage[node_id.idx] - if iszero(s) - concentration_state[node_id.idx, :] .= 0 - else - concentration_state[node_id.idx, :] .= - mass[node_id.idx] ./ current_storage[node_id.idx] + infiltrated_volume = (u.infiltration[node_id.idx] - uprev.infiltration[node_id.idx]) / nsubsteps[node_id.idx] + mass_node .-= concentration_state[node_id.idx, :] .* infiltrated_volume + + # Take care of infinitely small masses, possibly becoming negative due to truncation. + for I in eachindex(mass_node) + if (-eps(Float64)) < mass_node[I] < (eps(Float64)) + mass_node[I] = 0.0 + end + end + + # Check for negative masses + if any(<(0), mass_node) + errors = true + for substance_idx in findall(<(0), mass_node) + substance_name = basin.concentration_data.substances[substance_idx] + substance_mass = mass_node[substance_idx] + @error "$node_id has negative mass $substance_mass for substance $substance_name" + end + end + + # Update the Basin concentrations again based on the removed mass + s = basin.storage_prev[node_id.idx] + (current_storage[node_id.idx] - basin.storage_prev[node_id.idx]) * substep / max_substeps + if iszero(s) + concentration_state[node_id.idx, :] .= 0 + else + concentration_state[node_id.idx, :] .= + mass[node_id.idx] ./ s + end + + # Debug: check continuity tracer after outflow step + c_cont = concentration_state[node_id.idx, Substance.Continuity] + if !iszero(s) && abs(c_cont - 1.0) > 0.01 + @warn "Continuity drift after outflows" node_id substep c_cont s mass_cont = mass[node_id.idx][Substance.Continuity] + end end + errors && error("Negative mass(es) detected at t = $(tprev + substep * dt_sub) s") end - errors && error("Negative mass(es) detected at t = $t s") - basin.storage_prev .= current_storage basin.level_prev .= current_level return nothing @@ -558,8 +594,13 @@ function check_water_balance_error!( end function save_solver_stats(u, t, integrator) - (; dt) = integrator + (; dt, p) = integrator (; stats) = integrator.sol + (; p_independent) = p + (; basin) = p_independent + (; concentration_data) = basin + (; nsubsteps) = concentration_data + return (; time = t, time_ns = time_ns(), @@ -567,6 +608,7 @@ function save_solver_stats(u, t, integrator) linear_solves = stats.nsolve, accepted_timesteps = stats.naccept, rejected_timesteps = stats.nreject, + max_subtimesteps = maximum(nsubsteps), dt, ) end diff --git a/core/src/concentration.jl b/core/src/concentration.jl index 7c2dee74be..7a4902fdee 100644 --- a/core/src/concentration.jl +++ b/core/src/concentration.jl @@ -2,10 +2,10 @@ Process mass inflows from UserDemand separately as the inflow and outflow are decoupled in the states """ -function mass_inflows_from_user_demand!(integrator::DEIntegrator)::Nothing - (; p, t) = integrator +function mass_inflows_from_user_demand!(integrator::DEIntegrator, substep = 1, max_substeps = 1)::Nothing + (; p, tprev, dt) = integrator (; basin, user_demand) = p.p_independent - (; concentration_state, mass) = basin.concentration_data + (; concentration_state, mass, nsubsteps, stepsize) = basin.concentration_data for (inflow_link, outflow_link) in zip(user_demand.inflow_link, user_demand.outflow_link) @@ -13,26 +13,32 @@ function mass_inflows_from_user_demand!(integrator::DEIntegrator)::Nothing to_node = outflow_link.link[2] user_demand_idx = outflow_link.link[1].idx - cumulative_user_demand_outflow = flow_update_on_link(integrator, outflow_link.link) + if to_node.type == NodeType.Basin && (substep % stepsize[to_node.idx]) == 0 if to_node.type == NodeType.Basin # Pass through all upstream substance concentrations. # Note that when outflow < inflow UserDemand consumes the # difference including the substances. + # Exclude the UserDemand tracer from upstream: save before, # restore after, so only the fresh tracer from # add_substance_mass! (= 1.0) ends up in the return flow. ud_mass_before = mass[to_node.idx][Substance.UserDemand] + + # Substance added from upstream + # Note that when outflow < inflow UserDemand consumes the difference including the substances + cumulative_user_demand_outflow = flow_update_on_link(integrator, outflow_link.link) mass[to_node.idx] .+= - concentration_state[from_node.idx, :] .* cumulative_user_demand_outflow + concentration_state[from_node.idx, :] .* cumulative_user_demand_outflow / nsubsteps[to_node.idx] + mass[to_node.idx][Substance.UserDemand] = ud_mass_before # Add fresh UserDemand tracer (= 1.0) and any user-defined substances add_substance_mass!( mass[to_node.idx], user_demand.concentration_itp[user_demand_idx], - cumulative_user_demand_outflow, - t, + cumulative_user_demand_outflow / nsubsteps[to_node.idx], + tprev + dt / max_substeps * substep, ) end end @@ -42,10 +48,10 @@ end """ Process all mass inflows to basins """ -function mass_inflows_basin!(integrator::DEIntegrator)::Nothing - (; p, t) = integrator +function mass_inflows_basin!(integrator::DEIntegrator, substep = 1, max_substeps = 1)::Nothing + (; p, tprev, dt) = integrator (; basin, state_inflow_link, state_outflow_link, level_boundary) = p.p_independent - (; cumulative_in, concentration_state, mass) = basin.concentration_data + (; cumulative_in, concentration_state, mass, nsubsteps, stepsize) = basin.concentration_data # Loop over connections that have state for (inflow_link, outflow_link) in zip(state_inflow_link, state_outflow_link) @@ -58,20 +64,21 @@ function mass_inflows_basin!(integrator::DEIntegrator)::Nothing continue end - if from_node.type == NodeType.Basin + if from_node.type == NodeType.Basin && (substep % stepsize[from_node.idx]) == 0 + cumulative_flow = flow_update_on_link(integrator, inflow_link.link) # Negative flow over the inflow link means flow into the from_node if cumulative_flow < 0 - cumulative_in[from_node.idx] -= cumulative_flow + cumulative_in[from_node.idx] -= cumulative_flow / nsubsteps[from_node.idx] if to_node.type == NodeType.Basin mass[from_node.idx] .-= - concentration_state[to_node.idx, :] .* cumulative_flow + concentration_state[to_node.idx, :] .* cumulative_flow / nsubsteps[from_node.idx] elseif to_node.type == NodeType.LevelBoundary add_substance_mass!( mass[from_node.idx], level_boundary.concentration_itp[to_node.idx], - -cumulative_flow, - t, + -cumulative_flow / nsubsteps[from_node.idx], + tprev + dt / max_substeps * substep, ) elseif (to_node.type == NodeType.Terminal && to_node.value == 0) # UserDemand inflow is discoupled from its outflow @@ -83,20 +90,21 @@ function mass_inflows_basin!(integrator::DEIntegrator)::Nothing end end - if to_node.type == NodeType.Basin + if to_node.type == NodeType.Basin && (substep % stepsize[to_node.idx]) == 0 + cumulative_flow = flow_update_on_link(integrator, outflow_link.link) if cumulative_flow > 0 - cumulative_in[to_node.idx] += cumulative_flow + cumulative_in[to_node.idx] += cumulative_flow / nsubsteps[to_node.idx] if from_node.type == NodeType.Basin mass[to_node.idx] .+= - concentration_state[from_node.idx, :] .* cumulative_flow + concentration_state[from_node.idx, :] .* cumulative_flow / nsubsteps[to_node.idx] elseif from_node.type == NodeType.LevelBoundary add_substance_mass!( mass[to_node.idx], level_boundary.concentration_itp[from_node.idx], - cumulative_flow, - t, + cumulative_flow / nsubsteps[to_node.idx], + tprev + dt / max_substeps * substep, ) elseif from_node.type == NodeType.Terminal && from_node.value == 0 # The unset flow link defaults to Terminal #0 @@ -113,24 +121,24 @@ end """ Process all mass outflows from Basins """ -function mass_outflows_basin!(integrator::DEIntegrator)::Nothing +function mass_outflows_basin!(integrator::DEIntegrator, substep = 1, max_substeps = 1)::Nothing (; state_inflow_link, state_outflow_link, basin) = integrator.p.p_independent - (; mass, concentration_state) = basin.concentration_data + (; mass, concentration_state, nsubsteps, stepsize) = basin.concentration_data @views for (inflow_link, outflow_link) in zip(state_inflow_link, state_outflow_link) from_node = inflow_link.link[1] to_node = outflow_link.link[2] - if from_node.type == NodeType.Basin + if from_node.type == NodeType.Basin && (substep % stepsize[from_node.idx]) == 0 flow = flow_update_on_link(integrator, inflow_link.link) if flow > 0 - mass[from_node.idx] .-= concentration_state[from_node.idx, :] .* flow + mass[from_node.idx] .-= concentration_state[from_node.idx, :] .* flow / nsubsteps[from_node.idx] end end - if to_node.type == NodeType.Basin + if to_node.type == NodeType.Basin && (substep % stepsize[to_node.idx]) == 0 flow = flow_update_on_link(integrator, outflow_link.link) if flow < 0 - mass[to_node.idx] .+= concentration_state[to_node.idx, :] .* flow + mass[to_node.idx] .+= concentration_state[to_node.idx, :] .* flow / nsubsteps[to_node.idx] end end end diff --git a/core/src/config.jl b/core/src/config.jl index 166c90d3a5..a49555f135 100644 --- a/core/src/config.jl +++ b/core/src/config.jl @@ -187,7 +187,6 @@ end maxiters::Int = 1.0e9 sparse::Bool = true autodiff::Bool = true - evaporate_mass::Bool = true depth_threshold::Float64 = 0.1 level_difference_threshold::Float64 = 0.02 specialize::Bool = false @@ -223,6 +222,12 @@ end route_priority::RoutePriority = RoutePriority() end +@option struct Concentration <: TableOption + evaporate_mass::Bool = true + substep_ratio::Float64 = 1 / 5 + substep_depth::Int = 10 +end + @option struct Experimental <: TableOption concentration::Bool = false allocation::Bool = false @@ -248,6 +253,7 @@ end results_dir::String interpolation::Interpolation = Interpolation() allocation::Allocation = Allocation() + concentration::Concentration = Concentration() solver::Solver = Solver() logging::Logging = Logging() results::Results = Results() diff --git a/core/src/parameter.jl b/core/src/parameter.jl index a5c5263640..fc9e1ad9aa 100644 --- a/core/src/parameter.jl +++ b/core/src/parameter.jl @@ -5,6 +5,7 @@ const SolverStats = @NamedTuple{ linear_solves::Int, accepted_timesteps::Int, rejected_timesteps::Int, + max_subtimesteps::Int, dt::Float64, } @@ -420,8 +421,16 @@ abstract type AbstractDemandNode <: AbstractParameterNode end @kwdef struct ConcentrationData # Config setting to enable/disable evaporation of mass evaporate_mass::Bool = true + # Safety factor to determine number of substeps based on residence time + substep_ratio::Float64 = 1 / 5 + # Maximum number of substeps (2^substep_depth) + substep_depth::Int = 10 # Cumulative inflow for each Basin at a given time cumulative_in::Vector{Float64} = zeros(Float64, 0) + # Number of substeps per Basin for mass outflow processing + nsubsteps::Vector{Int16} = zeros(Int16, 0) + # Stepsize for each Basin for mass outflow processing, in number of substeps (e.g. stepsize of 2 means every other substep is processed) + stepsize::Vector{Int16} = zeros(Int16, 0) # Matrix with concentrations for each Basin and substance concentration_state::Matrix{Float64} = zeros(Float64, 0, 0) # Basin, substance # Vectors with concentration timeseries interpolations for each incoming forcing per Basin per substance diff --git a/core/src/read.jl b/core/src/read.jl index cd3a7f2b8f..b589612140 100644 --- a/core/src/read.jl +++ b/core/src/read.jl @@ -802,9 +802,13 @@ function ConcentrationData( end cumulative_in = zeros(n_basin) + nsubsteps = fill(one(Int16), n_basin) + stepsize = fill(one(Int16), n_basin) return ConcentrationData(; - config.solver.evaporate_mass, + config.concentration.evaporate_mass, + config.concentration.substep_ratio, + config.concentration.substep_depth, concentration_state, concentration_itp_drainage, concentration_itp_precipitation, @@ -814,6 +818,8 @@ function ConcentrationData( concentration_external, substances, cumulative_in, + nsubsteps, + stepsize, ) end diff --git a/core/src/write.jl b/core/src/write.jl index 1cbfb578a7..bbe6ff26dd 100644 --- a/core/src/write.jl +++ b/core/src/write.jl @@ -275,6 +275,8 @@ const CF = OrderedDict{String, OrderedDict{String, String}}( OrderedDict("units" => "1", "long_name" => "number of accepted timesteps"), "rejected_timesteps" => OrderedDict("units" => "1", "long_name" => "number of rejected timesteps"), + "max_subtimesteps" => + OrderedDict("units" => "1", "long_name" => "maximum number of subtimesteps in the last timestep"), "dt" => OrderedDict("units" => "s", "long_name" => "timestep size"), "from_node_id" => OrderedDict("long_name" => "source node identifier"), "to_node_id" => OrderedDict("long_name" => "destination node identifier"), @@ -434,6 +436,7 @@ function solver_stats_data(model::Model; table::Bool = true) linear_solves = diff(solver_stats.linear_solves), accepted_timesteps = diff(solver_stats.accepted_timesteps), rejected_timesteps = diff(solver_stats.rejected_timesteps), + max_subtimesteps = solver_stats.max_subtimesteps[2:end], dt = solver_stats.dt[2:end], ) end From 06d71e4465ca9db083131aac72af3764d58d6ebf Mon Sep 17 00:00:00 2001 From: Maarten Pronk Date: Wed, 1 Apr 2026 21:36:30 +0200 Subject: [PATCH 2/2] Wrong merge. --- core/src/concentration.jl | 1 - 1 file changed, 1 deletion(-) diff --git a/core/src/concentration.jl b/core/src/concentration.jl index 7a4902fdee..76c8426233 100644 --- a/core/src/concentration.jl +++ b/core/src/concentration.jl @@ -15,7 +15,6 @@ function mass_inflows_from_user_demand!(integrator::DEIntegrator, substep = 1, m if to_node.type == NodeType.Basin && (substep % stepsize[to_node.idx]) == 0 - if to_node.type == NodeType.Basin # Pass through all upstream substance concentrations. # Note that when outflow < inflow UserDemand consumes the # difference including the substances.