From c0349ba8e3f0bcc55d8f76d5c01b56018fe5ac8c Mon Sep 17 00:00:00 2001 From: Bart de Koning <74617371+SouthEndMusic@users.noreply.github.com> Date: Fri, 28 Jul 2023 08:48:45 +0200 Subject: [PATCH] Jacobian prototype (sparsity) (#459) Fixes part of https://github.com/Deltares/Ribasim/issues/419. --------- Co-authored-by: Bart de Koning Co-authored-by: Martijn Visser --- core/src/bmi.jl | 6 +- core/src/utils.jl | 193 +++++++++++++++++++++++++++++++++++- core/src/validation.jl | 8 +- core/test/utils.jl | 33 ++++++ docs/contribute/addnode.qmd | 2 + 5 files changed, 235 insertions(+), 7 deletions(-) diff --git a/core/src/bmi.jl b/core/src/bmi.jl index ae907409c..c07129f56 100644 --- a/core/src/bmi.jl +++ b/core/src/bmi.jl @@ -54,8 +54,12 @@ function BMI.initialize(T::Type{Model}, config::Config)::Model # for Float32 this method allows max ~1000 year simulations without accuracy issues @assert eps(t_end) < 3600 "Simulation time too long" timespan = (zero(t_end), t_end) + + jac_prototype = get_jac_prototype(parameters) + RHS = ODEFunction(water_balance!; jac_prototype) + @timeit_debug to "Setup ODEProblem" begin - prob = ODEProblem(water_balance!, u0, timespan, parameters) + prob = ODEProblem(RHS, u0, timespan, parameters) end callback, saved_flow = create_callbacks(parameters) diff --git a/core/src/utils.jl b/core/src/utils.jl index 100cbc42e..cc209ff67 100644 --- a/core/src/utils.jl +++ b/core/src/utils.jl @@ -352,8 +352,9 @@ function expand_logic_mapping( for (node_id, truth_state) in keys(logic_mapping) pattern = r"^[TF\*]+$" - msg = "Truth state \'$truth_state\' contains illegal characters or is empty." - @assert occursin(pattern, truth_state) msg + if !occursin(pattern, truth_state) + error("Truth state \'$truth_state\' contains illegal characters or is empty.") + end control_state = logic_mapping[(node_id, truth_state)] n_wildcards = count(==('*'), truth_state) @@ -391,3 +392,191 @@ function expand_logic_mapping( end return logic_mapping_expanded end + +"""Get all node fieldnames of the parameter object.""" +nodefields(p::Parameters) = ( + name for + name in fieldnames(typeof(p)) if fieldtype(typeof(p), name) <: AbstractParameterNode +) + +""" +Get a sparse matrix whose sparsity matches the sparsity of the Jacobian +of the ODE problem. All nodes are taken into consideration, also the ones +that are inactive. +""" +function get_jac_prototype(p::Parameters)::SparseMatrixCSC{Float64, Int64} + (; basin, pid_control) = p + + n_basins = length(basin.node_id) + n_states = n_basins + length(pid_control.node_id) + jac_prototype = spzeros(n_states, n_states) + + for nodefield in nodefields(p) + update_jac_prototype!(jac_prototype, p, getfield(p, nodefield)) + end + + return jac_prototype +end + +""" +If both the unique node upstream and the unique node downstream of these +nodes are basins, then these directly depend on eachother and affect the Jacobian 2x +""" +function update_jac_prototype!( + jac_prototype::SparseMatrixCSC{Float64, Int64}, + p::Parameters, + node::Union{LinearResistance, ManningResistance}, +)::Nothing + (; basin, connectivity) = p + (; graph_flow) = connectivity + + for id in node.node_id + + # Only if the inneighbor and the outneighbor are basins + # do we get a contribution + id_in = only(inneighbors(graph_flow, id)) + id_out = only(outneighbors(graph_flow, id)) + + has_index_in, idx_in = id_index(basin.node_id, id_in) + has_index_out, idx_out = id_index(basin.node_id, id_out) + + if has_index_in && has_index_out + jac_prototype[idx_in, idx_out] = 1.0 + jac_prototype[idx_out, idx_in] = 1.0 + end + end + return nothing +end + +""" +Method for nodes that do not contribute to the Jacobian +""" +function update_jac_prototype!( + jac_prototype::SparseMatrixCSC{Float64, Int64}, + p::Parameters, + node::AbstractParameterNode, +)::Nothing + node_type = nameof(typeof(node)) + + if !isa( + node, + Union{ + Basin, + DiscreteControl, + FlowBoundary, + FractionalFlow, + LevelBoundary, + Terminal, + }, + ) + error( + "It is not specified how nodes of type $node_type contribute to the Jacobian prototype.", + ) + end +end + +""" +If both the unique node upstream and the nodes down stream (or one node further +if a fractional flow is in between) are basins, then the downstream basin depends +on the upstream basin(s) and affect the Jacobian as many times as there are downstream basins +""" +function update_jac_prototype!( + jac_prototype::SparseMatrixCSC{Float64, Int64}, + p::Parameters, + node::Union{Pump, TabulatedRatingCurve}, +)::Nothing + (; basin, fractional_flow, connectivity) = p + (; graph_flow) = connectivity + + for id in node.node_id + id_in = only(inneighbors(graph_flow, id)) + + # For inneighbors only directly connected basins give a contribution + has_index_in, idx_in = id_index(basin.node_id, id_in) + + # For outneighbors there can be directly connected basins + # or basins connected via a fractional flow + if has_index_in + idxs_out = + get_fractional_flow_connected_basins(id, basin, fractional_flow, graph_flow) + + if isempty(idxs_out) + id_out = only(outneighbors(graph_flow, id)) + has_index_out, idx_out = id_index(basin.node_id, id_out) + + if has_index_out + push!(idxs_out, idx_out) + end + end + + for idx_out in idxs_out + jac_prototype[idx_in, idx_out] = 1.0 + end + end + end + return nothing +end + +""" +The controlled basin affects itself and the basins upstream and downstream of the controlled pump +affect eachother if there is a basin upstream of the pump. The state for the integral term +and the controlled basin affect eachother, and the same for the integral state and the basin +upstream of the pump if it is indeed a basin. +""" +function update_jac_prototype!( + jac_prototype::SparseMatrixCSC{Float64, Int64}, + p::Parameters, + node::PidControl, +)::Nothing + (; basin, connectivity) = p + (; graph_control) = connectivity + + n_basins = length(basin.node_id) + + for (pid_idx, (listen_node_id, id)) in enumerate(zip(node.listen_node_id, node.node_id)) + id_out = only(outneighbors(graph_control, id)) + id_out_in = only(inneighbors(graph_control, id_out)) + + _, listen_idx = id_index(basin.node_id, listen_node_id) + + # PID control integral state + pid_state_idx = n_basins + pid_idx + jac_prototype[pid_state_idx, listen_idx] = 1.0 + jac_prototype[listen_idx, pid_state_idx] = 1.0 + + # The basin upstream of the pump + has_index, idx_out_in = id_index(basin.node_id, id_out_in) + + if has_index + jac_prototype[pid_state_idx, idx_out_in] = 1.0 + jac_prototype[idx_out_in, pid_state_idx] = 1.0 + + # The basin upstream of the pump also depends on the controlled basin + jac_prototype[listen_idx, idx_out_in] = 1.0 + end + end + return nothing +end + +""" +Get the state index of the basins that are connected to a node of given id via fractional flow. +""" +function get_fractional_flow_connected_basins( + node_id::Int, + basin::Basin, + fractional_flow::FractionalFlow, + graph_flow::DiGraph{Int}, +)::Vector{Int} + basin_idxs = Int[] + + for first_outneighbor_id in outneighbors(graph_flow, node_id) + if first_outneighbor_id in fractional_flow.node_id + second_outneighbor_id = only(outneighbors(graph_flow, first_outneighbor_id)) + has_index, basin_idx = id_index(basin.node_id, second_outneighbor_id) + if has_index + push!(basin_idxs, basin_idx) + end + end + end + return basin_idxs +end diff --git a/core/src/validation.jl b/core/src/validation.jl index 12664b0f2..4e00149c1 100644 --- a/core/src/validation.jl +++ b/core/src/validation.jl @@ -50,8 +50,7 @@ neighbortypes(::Val{:Basin}) = Set(( :FlowBoundary, )) neighbortypes(::Val{:Terminal}) = Set{Symbol}() # only endnode -neighbortypes(::Val{:FractionalFlow}) = - Set((:Basin, :FractionalFlow, :Terminal, :LevelBoundary)) +neighbortypes(::Val{:FractionalFlow}) = Set((:Basin, :Terminal, :LevelBoundary)) neighbortypes(::Val{:FlowBoundary}) = Set((:Basin, :FractionalFlow, :Terminal, :LevelBoundary)) neighbortypes(::Val{:LevelBoundary}) = Set((:LinearResistance, :ManningResistance, :Pump)) @@ -293,8 +292,9 @@ function sorted_table!( by = sort_by_function(table) if Tables.getcolumn(table, :node_id) isa Arrow.Primitive et = eltype(table) - msg = "Arrow table for $et not sorted as required." - @assert issorted(table; by) msg + if !issorted(table; by) + error("Arrow table for $et not sorted as required.") + end else sort!(table; by) end diff --git a/core/test/utils.jl b/core/test/utils.jl index 64b61c2f0..e448beb39 100644 --- a/core/test/utils.jl +++ b/core/test/utils.jl @@ -3,6 +3,7 @@ using Dictionaries: Indices using Test using DataInterpolations: LinearInterpolation using StructArrays: StructVector +using SQLite @testset "id_index" begin ids = Indices([2, 4, 6]) @@ -99,3 +100,35 @@ end logic_mapping, ) end + +@testset "Jacobian sparsity" begin + toml_path = normpath(@__DIR__, "../../data/basic/basic.toml") + + cfg = Ribasim.parsefile(toml_path) + gpkg_path = Ribasim.input_path(cfg, cfg.geopackage) + db = SQLite.DB(gpkg_path) + + p = Ribasim.Parameters(db, cfg) + jac_prototype = Ribasim.get_jac_prototype(p) + + @test jac_prototype.m == 4 + @test jac_prototype.n == 4 + @test jac_prototype.colptr == [1, 2, 3, 4, 6] + @test jac_prototype.rowval == [2, 1, 2, 2, 3] + @test jac_prototype.nzval == ones(5) + + toml_path = normpath(@__DIR__, "../../data/pid_1/pid_1.toml") + + cfg = Ribasim.parsefile(toml_path) + gpkg_path = Ribasim.input_path(cfg, cfg.geopackage) + db = SQLite.DB(gpkg_path) + + p = Ribasim.Parameters(db, cfg) + jac_prototype = Ribasim.get_jac_prototype(p) + + @test jac_prototype.m == 2 + @test jac_prototype.n == 2 + @test jac_prototype.colptr == [1, 2, 3] + @test jac_prototype.rowval == [2, 1] + @test jac_prototype.nzval == ones(2) +end diff --git a/docs/contribute/addnode.qmd b/docs/contribute/addnode.qmd index f696c827e..4dace81a5 100644 --- a/docs/contribute/addnode.qmd +++ b/docs/contribute/addnode.qmd @@ -70,6 +70,8 @@ function formulate!(new_node_type::NewNodeType, p::Parameters)::Nothing end ``` +If the new node type can create dependendies between basins, this has to be taken into account in the sparsity structure of the Jacobian of `water_balance!` which is constructed in the function `get_jac_prototype` in `src\utils.jl`. + # Python I/O