diff --git a/core/src/Ribasim.jl b/core/src/Ribasim.jl index ac53f04cd..c6ec1b35c 100644 --- a/core/src/Ribasim.jl +++ b/core/src/Ribasim.jl @@ -67,7 +67,7 @@ using SciMLBase: ODEFunction, ODEProblem, ODESolution, - VectorContinuousCallback, + DiscreteCallback, get_proposed_dt using SparseArrays: SparseMatrixCSC, spzeros using SQLite: SQLite, DB, Query, esc_id diff --git a/core/src/callback.jl b/core/src/callback.jl index aba885d5c..16f8ab422 100644 --- a/core/src/callback.jl +++ b/core/src/callback.jl @@ -21,11 +21,11 @@ function TrapezoidIntegrationCallback(integrand_func!, integral_value)::Discrete end function (affect!::TrapezoidIntegrationAffect)(integrator)::Nothing - (; integrand_func, integrand_value, integrand_value_prev, cache, integral) = affect! - (; dt) = integrand + (; integrand_func!, integrand_value, integrand_value_prev, cache, integral) = affect! + (; dt) = integrator copyto!(integrand_value_prev, integrand_value) - integrand_func(integrand_value, integrator.p) + integrand_func!(integrand_value, integrator.p) cache += integrand_value_prev cache += integrand_value @@ -65,6 +65,11 @@ function create_callbacks( TrapezoidIntegrationCallback(flow_integrand!, graph[].integrated_flow) push!(callbacks, integrating_flow_cb) + # Save mean flows + saved_flow = SavedValues(Float64, SavedFlow) + save_flow_cb = SavingCallback(save_flow, saved_flow; saveat, save_start = false) + push!(callbacks, save_flow_cb) + # Integrate vertical basin fluxes for BMI integrated_bmi_data = ComponentVector{Float64}(; get_tmp(basin.vertical_flux, 0)..., @@ -124,7 +129,6 @@ function flow_integrand!(out, p)::Nothing out.flow .= flow vertical_flux_view(out) .= vertical_flux - set_inoutflows!(out, graph, basin) return nothing end @@ -153,41 +157,14 @@ function allocation_integrand!(out, p)::Nothing 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[] - (; node_id) = integrator.p.basin - + (; basin, graph) = integrator.p + (; integrated_flow) = graph[] Δt = get_Δt(integrator) - flow_mean = copy(flow_integrated) + flow_mean = copy(integrated_flow) flow_mean ./= Δt - fill!(flow_integrated, 0.0) - - # Divide the flows over edges to Basin inflow and outflow, regardless of edge direction. - inflow_mean = zeros(length(node_id)) - outflow_mean = zeros(length(node_id)) - - for (i, basin_id) in enumerate(node_id) - for inflow_id in inflow_ids(graph, basin_id) - q = flow_mean[flow_dict[inflow_id, basin_id]] - if q > 0 - inflow_mean[i] += q - else - outflow_mean[i] -= q - end - end - for outflow_id in outflow_ids(graph, basin_id) - q = flow_mean[flow_dict[basin_id, outflow_id]] - if q > 0 - outflow_mean[i] += q - else - inflow_mean[i] -= q - end - end - end - + fill!(integrated_flow, 0.0) + inflow_mean, outflow_mean = compute_mean_inoutflows(flow_mean.flow, graph, basin) return SavedFlow(; flow = flow_mean, inflow = inflow_mean, outflow = outflow_mean) end diff --git a/core/src/graph.jl b/core/src/graph.jl index ea39611a1..e73533a4b 100644 --- a/core/src/graph.jl +++ b/core/src/graph.jl @@ -104,11 +104,19 @@ function create_graph(db::DB, config::Config, chunk_sizes::Vector{Int})::MetaGra flow = DiffCache(flow, chunk_sizes) end - flow = zero(flow_counter) + flow = zeros(flow_counter) n_basins = length(get_ids(db, "Basin")) - vertical_flux = zero(n_basins) - integrated_flow = ComponentVector{Float64}(; flow, vertical_flux...) + integrated_flow = ComponentVector{Float64}(; + flow, + precipitation = zeros(n_basins), + evaporation = zeros(n_basins), + drainage = zeros(n_basins), + infiltration = zeros(n_basins), + ) + if config.solver.autodiff + flow = DiffCache(flow, chunk_sizes) + end graph_data = (; node_ids, edges_source, flow_dict, flow, integrated_flow, config.solver.saveat) graph = @set graph.graph_data = graph_data diff --git a/core/src/model.jl b/core/src/model.jl index 94f5aefd3..5aef1ab81 100644 --- a/core/src/model.jl +++ b/core/src/model.jl @@ -1,7 +1,7 @@ -struct SavedResults{V1 <: ComponentVector{Float64}, V2 <: ComponentVector{Float64}} - flow::IntegrandValues{Float64, V1} +struct SavedResults{V1 <: ComponentVector{Float64}} + flow::SavedValues{Float64, SavedFlow} subgrid_level::SavedValues{Float64, Vector{Float64}} - integrated_bmi_data::V2 + integrated_bmi_data::V1 end """ diff --git a/core/src/parameter.jl b/core/src/parameter.jl index 9503ce76c..4fa98fd1a 100644 --- a/core/src/parameter.jl +++ b/core/src/parameter.jl @@ -136,6 +136,19 @@ end abstract type AbstractParameterNode end +""" +In-memory storage of saved mean flows for writing to results. + +- `flow`: The mean flows on all edges +- `inflow`: The sum of the mean flows coming into each basin +- `outflow`: The sum of the mean flows going out of each basin +""" +@kwdef struct SavedFlow{V1 <: ComponentVector{Float64}} + flow::V1 + inflow::Vector{Float64} + outflow::Vector{Float64} +end + """ Requirements: diff --git a/core/src/read.jl b/core/src/read.jl index d6f0584a8..1020d2d56 100644 --- a/core/src/read.jl +++ b/core/src/read.jl @@ -1035,7 +1035,7 @@ function Allocation(db::DB, config::Config, graph::MetaGraph)::Allocation integrated_flow_mapping[(node_id, node_id)] = flow_counter end end - integrated_flow = IntegrandValuesSum(zeros(flow_counter)) + integrated_flow = zeros(flow_counter) return Allocation( Int32[], diff --git a/core/src/util.jl b/core/src/util.jl index 54b0ef561..f48d3c7ae 100644 --- a/core/src/util.jl +++ b/core/src/util.jl @@ -521,6 +521,9 @@ function Base.getindex(fv::FlatVector, i::Int) return v[r + 1] end +"Construct a FlatVector from one of the fields of SavedFlow." +FlatVector(saveval::Vector{SavedFlow}, sym::Symbol) = FlatVector(getfield.(saveval, sym)) + """ Function that goes smoothly from 0 to 1 in the interval [0,threshold], and is constant outside this interval. @@ -758,28 +761,34 @@ function vertical_flux_view(v::ComponentVector) return @view v[(:precipitation, :evaporation, :drainage, :infiltration)] end -function set_inoutflows!( - v::ComponentVector{Float64}, +function compute_mean_inoutflows( + flow_mean::AbstractVector, graph::MetaGraph, basin::Basin, -)::Nothing - for (i, basin_id) in enumerate(basin.node_id) +)::Tuple{Vector{Float64}, Vector{Float64}} + (; node_id) = basin + + # Divide the flows over edges to Basin inflow and outflow, regardless of edge direction. + inflow_mean = zeros(length(node_id)) + outflow_mean = zeros(length(node_id)) + + for (i, basin_id) in enumerate(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 + inflow_mean[i] += q else - v.basin_outflow[i] -= q + outflow_mean[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 + inflow_mean[i] += q else - v.basin_inflow[i] -= q + outflow_mean[i] -= q end end end - return nothing + return inflow_mean, outflow_mean end diff --git a/core/src/write.jl b/core/src/write.jl index b1af36ee1..2002c7d6e 100644 --- a/core/src/write.jl +++ b/core/src/write.jl @@ -94,7 +94,7 @@ function basin_table( balance_error::Vector{Float64}, relative_error::Vector{Float64}, } - (; saved, integrator) = model + (; saved) = model # The last timestep is not included; there is no period over which to compute flows. data = get_storages_and_levels(model) storage = vec(data.storage[:, begin:(end - 1)]) @@ -104,17 +104,37 @@ function basin_table( nbasin = length(data.node_id) ntsteps = length(data.time) - 1 nrows = nbasin * ntsteps - saveats = integrator.sol.t - - inflow_rate = average_over_saveats(saved.flow, :basin_inflow, saveats) - outflow_rate = average_over_saveats(saved.flow, :basin_inflow, saveats) - precipitation = average_over_saveats(saved.flow, :precipitation, saveats) - evaporation = average_over_saveats(saved.flow, :evaporation, saveats) - drainage = average_over_saveats(saved.flow, :drainage, saveats) - infiltration = average_over_saveats(saved.flow, :infiltration, saveats) + + inflow_rate = FlatVector(saved.flow.saveval, :inflow) + outflow_rate = FlatVector(saved.flow.saveval, :outflow) + precipitation = zeros(nrows) + evaporation = zeros(nrows) + drainage = zeros(nrows) + infiltration = zeros(nrows) balance_error = zeros(nrows) relative_error = zeros(nrows) + @show length(model.integrator.sol.t) + + idx_row = 0 + for cvec in saved.flow.saveval + for (precipitation_, evaporation_, drainage_, infiltration_) in zip( + cvec.flow.precipitation, + cvec.flow.evaporation, + cvec.flow.drainage, + cvec.flow.infiltration, + ) + idx_row += 1 + precipitation[idx_row] = precipitation_ + evaporation[idx_row] = evaporation_ + drainage[idx_row] = drainage_ + infiltration[idx_row] = infiltration_ + end + end + + @show ntsteps + @show length(saved.flow.saveval) + time = repeat(data.time[begin:(end - 1)]; inner = nbasin) Δtime_seconds = seconds.(diff(data.time)) Δtime = repeat(Δtime_seconds; inner = nbasin)