Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

use FixedSizeDiffCache for flows #581

Merged
merged 2 commits into from
Sep 8, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion core/src/Ribasim.jl
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ using Legolas: Legolas, @schema, @version, validate, SchemaVersion, declared
using Logging: current_logger, min_enabled_level, with_logger
using LoggingExtras: EarlyFilteredLogger, LevelOverrideLogger
using OrdinaryDiffEq
using PreallocationTools: DiffCache, get_tmp
using PreallocationTools: DiffCache, FixedSizeDiffCache, get_tmp
using SciMLBase
using SparseArrays
using SQLite: SQLite, DB, Query, esc_id
Expand Down
5 changes: 3 additions & 2 deletions core/src/create.jl
Original file line number Diff line number Diff line change
Expand Up @@ -202,10 +202,11 @@ function Connectivity(db::DB, config::Config, chunk_size::Int)::Connectivity
edge_ids_flow_inv = Dictionary(values(edge_ids_flow), keys(edge_ids_flow))

flow = adjacency_matrix(graph_flow, Float64)
nonzeros(flow) .= 0.0
flow .= 0.0

if config.solver.autodiff
flow = DiffCache(flow, chunk_size)
# FixedSizeDiffCache performs better for sparse matrix
flow = FixedSizeDiffCache(flow, chunk_size)
end

return Connectivity(
Expand Down
80 changes: 45 additions & 35 deletions core/src/solve.jl
Original file line number Diff line number Diff line change
Expand Up @@ -401,21 +401,21 @@ struct PidControl{T} <: AbstractParameterNode
end

# TODO Automatically add all nodetypes here
struct Parameters
struct Parameters{T, TSparse, C1, C2}
starttime::DateTime
connectivity::Connectivity
basin::Basin
connectivity::Connectivity{TSparse}
basin::Basin{T, C1}
linear_resistance::LinearResistance
manning_resistance::ManningResistance
tabulated_rating_curve::TabulatedRatingCurve
tabulated_rating_curve::TabulatedRatingCurve{C2}
fractional_flow::FractionalFlow
level_boundary::LevelBoundary
flow_boundary::FlowBoundary
pump::Pump
outlet::Outlet
pump::Pump{T}
outlet::Outlet{T}
terminal::Terminal
discrete_control::DiscreteControl
pid_control::PidControl
pid_control::PidControl{T}
lookup::Dict{Int, Symbol}
end

Expand Down Expand Up @@ -508,8 +508,8 @@ function formulate!(
du::AbstractVector,
basin::Basin,
storage::AbstractVector,
current_area,
current_level,
current_area::AbstractVector,
current_level::AbstractVector,
t::Float64,
)::Nothing
for i in eachindex(storage)
Expand Down Expand Up @@ -553,13 +553,13 @@ end
function continuous_control!(
u::ComponentVector,
du::ComponentVector,
current_area,
current_area::AbstractVector,
pid_control::PidControl,
p::Parameters,
integral_value::SubArray,
current_level::AbstractVector,
flow::SparseMatrixCSC,
pid_error,
flow::AbstractMatrix,
pid_error::AbstractVector,
t::Float64,
)::Nothing
(; connectivity, pump, outlet, basin, fractional_flow) = p
Expand Down Expand Up @@ -700,11 +700,11 @@ end
"""
Directed graph: outflow is positive!
"""
function formulate!(
function formulate_flow!(
linear_resistance::LinearResistance,
p::Parameters,
current_level,
flow,
current_level::AbstractVector,
flow::AbstractMatrix,
t::Float64,
)::Nothing
(; connectivity) = p
Expand Down Expand Up @@ -733,11 +733,11 @@ end
"""
Directed graph: outflow is positive!
"""
function formulate!(
function formulate_flow!(
tabulated_rating_curve::TabulatedRatingCurve,
p::Parameters,
current_level,
flow::SparseMatrixCSC,
current_level::AbstractVector,
flow::AbstractMatrix,
t::Float64,
)::Nothing
(; connectivity) = p
Expand Down Expand Up @@ -800,11 +800,11 @@ The average of the upstream and downstream water depth is used to compute cross-
hydraulic radius. This ensures that a basin can receive water after it has gone
dry.
"""
function formulate!(
function formulate_flow!(
manning_resistance::ManningResistance,
p::Parameters,
current_level,
flow,
current_level::AbstractVector,
flow::AbstractMatrix,
t::Float64,
)::Nothing
(; basin, connectivity) = p
Expand Down Expand Up @@ -859,7 +859,11 @@ function formulate!(
return nothing
end

function formulate!(fractional_flow::FractionalFlow, flow, p::Parameters)::Nothing
function formulate_flow!(
fractional_flow::FractionalFlow,
flow::AbstractMatrix,
p::Parameters,
)::Nothing
(; connectivity) = p
(; graph_flow) = connectivity
(; node_id, fraction) = fractional_flow
Expand All @@ -871,7 +875,12 @@ function formulate!(fractional_flow::FractionalFlow, flow, p::Parameters)::Nothi
return nothing
end

function formulate!(flow_boundary::FlowBoundary, p::Parameters, flow, t::Float64)::Nothing
function formulate_flow!(
flow_boundary::FlowBoundary,
p::Parameters,
flow::AbstractMatrix,
t::Float64,
)::Nothing
(; connectivity) = p
(; graph_flow) = connectivity
(; node_id, active, flow_rate) = flow_boundary
Expand All @@ -892,10 +901,10 @@ function formulate!(flow_boundary::FlowBoundary, p::Parameters, flow, t::Float64
end
end

function formulate!(
function formulate_flow!(
node::Union{Pump, Outlet},
p::Parameters,
flow,
flow::AbstractMatrix,
storage::AbstractVector,
)::Nothing
(; connectivity, basin) = p
Expand Down Expand Up @@ -938,7 +947,7 @@ end
function formulate!(
du::ComponentVector,
connectivity::Connectivity,
flow::SparseMatrixCSC,
flow::AbstractMatrix,
basin::Basin,
)::Nothing
# loop over basins
Expand All @@ -960,7 +969,7 @@ function formulate_flows!(
p::Parameters,
storage::AbstractVector,
current_level,
flow::SparseMatrixCSC,
flow::AbstractMatrix,
t::Float64,
)::Nothing
(;
Expand All @@ -973,13 +982,13 @@ function formulate_flows!(
outlet,
) = p

formulate!(linear_resistance, p, current_level, flow, t)
formulate!(manning_resistance, p, current_level, flow, t)
formulate!(tabulated_rating_curve, p, current_level, flow, t)
formulate!(flow_boundary, p, flow, t)
formulate!(fractional_flow, flow, p)
formulate!(pump, p, flow, storage)
formulate!(outlet, p, flow, storage)
formulate_flow!(linear_resistance, p, current_level, flow, t)
formulate_flow!(manning_resistance, p, current_level, flow, t)
formulate_flow!(tabulated_rating_curve, p, current_level, flow, t)
formulate_flow!(flow_boundary, p, flow, t)
formulate_flow!(fractional_flow, flow, p)
formulate_flow!(pump, p, flow, storage)
formulate_flow!(outlet, p, flow, storage)

return nothing
end
Expand All @@ -1000,7 +1009,8 @@ function water_balance!(

du .= 0.0
flow = get_tmp(connectivity.flow, u)
nonzeros(flow) .= 0.0
# use parent to avoid materializing the ReinterpretArray from FixedSizeDiffCache
parent(flow) .= 0.0

current_area = get_tmp(basin.current_area, u)
current_level = get_tmp(basin.current_level, u)
Expand Down
5 changes: 5 additions & 0 deletions core/test/run_models.jl
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,11 @@ end
end

@test model isa Ribasim.Model
p = model.integrator.p
@test p isa Ribasim.Parameters
@test isconcretetype(typeof(p))
@test all(isconcretetype, fieldtypes(typeof(p)))

@test successful_retcode(model)
@test model.integrator.sol.u[end] Float32[519.8817, 519.8798, 339.3959, 1418.4331] skip =
Sys.isapple() atol = 1.5
Expand Down
4 changes: 2 additions & 2 deletions docs/contribute/addnode.qmd
Original file line number Diff line number Diff line change
Expand Up @@ -61,10 +61,10 @@ end

## Node behavior

In general if the new node type dictates flow, the behaviour of the new node in the Ribasim core is defined in a method of the `formulate!` function, which is called within the `water_balance!` (both in `solve.jl`) function being the right hand side of the system of differential equations solved by Ribasim. Here the details depend highly on the specifics of the node type. An example structure of a `formulate!` method is given below.
In general if the new node type dictates flow, the behaviour of the new node in the Ribasim core is defined in a method of the `formulate_flow!` function, which is called within the `water_balance!` (both in `solve.jl`) function being the right hand side of the system of differential equations solved by Ribasim. Here the details depend highly on the specifics of the node type. An example structure of a `formulate_flow!` method is given below.

```julia
function formulate!(new_node_type::NewNodeType, p::Parameters)::Nothing
function formulate_flow!(new_node_type::NewNodeType, p::Parameters)::Nothing
# Retrieve relevant parameters
(; connectivity) = p
(; flow) = connectivity
Expand Down