Jacobian prototype (sparsity) (#459)

Fixes part of #419.

---------

Co-authored-by: Bart de Koning <[email protected]>
Co-authored-by: Martijn Visser <[email protected]>

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)

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

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

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

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