Update docstrings + documentation

core/src/allocation_optim.jl

@@ -256,7 +256,7 @@ function set_initial_capacities_source!(
 )::Nothing
     (; problem) = allocation_model
     (; graph, allocation) = p
-    (; flow_dict) = allocation
+    (; input_flow_dict) = allocation
     (; subnetwork_id) = allocation_model
     source_constraints = problem[:source]
     main_network_source_edges = get_main_network_connections(p, subnetwork_id)
@@ -267,7 +267,7 @@ function set_initial_capacities_source!(
             # If it is a source edge for this allocation problem
             if edge ∉ main_network_source_edges
                 # Reset the source to the averaged flow over the last allocation period
-                source_capacity = u.flow_allocation_input[flow_dict[edge]]
+                source_capacity = u.flow_allocation_input[input_flow_dict[edge]]
                 JuMP.set_normalized_rhs(
                     source_constraints[edge],
                     # It is assumed that the allocation procedure does not have to be differentiated.
@@ -362,11 +362,11 @@ function get_basin_data(
     (; graph, basin, level_demand, allocation) = p
     (; vertical_flux) = basin
     (; Δt_allocation) = allocation_model
-    (; flow_dict) = allocation
+    (; input_flow_dict) = allocation
     @assert node_id.type == NodeType.Basin
     vertical_flux = get_tmp(vertical_flux, 0)
     _, basin_idx = id_index(basin.node_id, node_id)
-    influx = u.flow_allocation_input[flow_dict[(node_id, node_id)]]
+    influx = u.flow_allocation_input[input_flow_dict[(node_id, node_id)]]
     _, basin_idx = id_index(basin.node_id, node_id)
     storage_basin = u.storage[basin_idx]
     control_inneighbors = inneighbor_labels_type(graph, node_id, EdgeType.control)

core/src/model.jl

@@ -35,6 +35,9 @@ function Model(config_path::AbstractString)::Model
     return Model(config)
 end
 
+"""
+Create the component state vector with the initial basin storages
+"""
 function initialize_state(db::DB, config::Config, basin::Basin)::ComponentVector
     n_states = get_n_states(db, config)
     u0 = ComponentVector{Float64}(
@@ -104,7 +107,8 @@ function Model(config::Config)::Model
 
         # initial state
         u0 = initialize_state(db, config, parameters.basin)
-        @assert length(u0.flow_allocation_input) == length(parameters.allocation.flow_dict) "Unexpected number of flows to integrate for allocation input."
+        @assert length(u0.flow_allocation_input) ==
+                length(parameters.allocation.input_flow_dict) "Unexpected number of flows to integrate for allocation input."
 
         sql = "SELECT node_id FROM Node ORDER BY node_id"
         node_id = only(execute(columntable, db, sql))

core/src/parameter.jl

@@ -64,7 +64,7 @@ main_network_connections: (from_id, to_id) from the main network to the subnetwo
 priorities: All used priority values.
 subnetwork_demands: The demand of an edge from the main network to a subnetwork
 subnetwork_allocateds: The allocated flow of an edge from the main network to a subnetwork
-flow_dict: ...
+input_flow_dict: Mapping from an edge (or basin forcings with a self loop) to an index in u.flow_allocation_input
 record_demand: A record of demands and allocated flows for nodes that have these
 record_flow: A record of all flows computed by allocation optimization, eventually saved to
     output file
@@ -76,7 +76,7 @@ struct Allocation
     priorities::Vector{Int32}
     subnetwork_demands::Dict{Tuple{NodeID, NodeID}, Vector{Float64}}
     subnetwork_allocateds::Dict{Tuple{NodeID, NodeID}, Vector{Float64}}
-    flow_dict::Dict{Tuple{NodeID, NodeID}, Int}
+    input_flow_dict::Dict{Tuple{NodeID, NodeID}, Int}
     record_demand::@NamedTuple{
         time::Vector{Float64},
         subnetwork_id::Vector{Int32},

core/src/read.jl

@@ -981,15 +981,15 @@ function Allocation(db::DB, config::Config, graph::MetaGraph)::Allocation
         optimization_type = String[],
     )
 
-    flow_dict = Dict{Tuple{NodeID, NodeID}, Int}()
+    input_flow_dict = Dict{Tuple{NodeID, NodeID}, Int}()
     flow_counter = 0
 
     # Find edges which serve as sources in allocation
     for edge_metadata in values(graph.edge_data)
         (; subnetwork_id_source, edge) = edge_metadata
         if subnetwork_id_source != 0
             flow_counter += 1
-            flow_dict[edge] = flow_counter
+            input_flow_dict[edge] = flow_counter
         end
     end
 
@@ -998,7 +998,7 @@ function Allocation(db::DB, config::Config, graph::MetaGraph)::Allocation
         if has_external_demand(graph, node_id, :level_demand)[1]
             edge = (node_id, node_id)
             flow_counter += 1
-            flow_dict[edge] = flow_counter
+            input_flow_dict[edge] = flow_counter
         end
     end
 
@@ -1009,7 +1009,7 @@ function Allocation(db::DB, config::Config, graph::MetaGraph)::Allocation
         get_all_priorities(db, config),
         Dict{Tuple{NodeID, NodeID}, Vector{Float64}}(),
         Dict{Tuple{NodeID, NodeID}, Vector{Float64}}(),
-        flow_dict,
+        input_flow_dict,
         record_demand,
         record_flow,
     )

core/src/solve.jl

@@ -614,8 +614,8 @@ function formulate_du_integration_flows!(
     forcings_integrated(du) .= vertical_flux
     forcings_bmi(du) .= vertical_flux
 
-    # Allocation_flows
-    for (edge, i) in allocation.flow_dict
+    # Allocation input flows
+    for (edge, i) in allocation.input_flow_dict
         du.flow_allocation_input[i] = if edge[1] == edge[2]
             get_influx(basin, edge[1])
         else

core/src/sparsity.jl

@@ -6,6 +6,13 @@ that are inactive.
 In Ribasim the Jacobian is typically sparse because each state only depends on a small
 number of other states.
 
+The aim is that jac_prototype[i,j] = 1.0 if and only if
+there is the possibility that at some point in the simulation:
+
+∂water_balance![j]/∂u[i] ≠ 0.0
+
+which means that du[j] depends on u[i].
+
 Note: the name 'prototype' does not mean this code is a prototype, it comes
 from the naming convention of this sparsity structure in the
 differentialequations.jl docs.

core/src/util.jl

@@ -766,6 +766,10 @@ function get_n_allocation_flow_inputs(db::DB)::Int
     return n_sources + n_level_demands
 end
 
+"""
+Get the number of states for each component of the
+state vector, and define the state names.
+"""
 function get_n_states(db::DB, config::Config)::NamedTuple
     n_basins = length(get_ids(db, "Basin"))
     n_pid_controls = length(get_ids(db, "PidControl"))
@@ -775,18 +779,25 @@ function get_n_states(db::DB, config::Config)::NamedTuple
         config.allocation.use_allocation ? get_n_allocation_flow_inputs(db) : 0
     # NOTE: This is the source of truth for the state component names
     return (;
+        # Basin storages
         storage = n_basins,
+        # PID control integral terms
         integral = n_pid_controls,
+        # Integrated flows for mean computation
         flow_integrated = n_flows,
+        # Integrated basin forcings for mean computation
         precipitation_integrated = n_basins,
         evaporation_integrated = n_basins,
         drainage_integrated = n_basins,
         infiltration_integrated = n_basins,
+        # Cumulative basin forcings for, for read or reset by BMI only
         precipitation_bmi = n_basins,
         evaporation_bmi = n_basins,
         drainage_bmi = n_basins,
         infiltration_bmi = n_basins,
+        # Flows averaged over Δt_allocation over edges that are allocation sources
         flow_allocation_input = n_allocation_flow_inputs,
+        # Cumulative UserDemand inflow volume, for read or reset by BMI only
         realized_user_demand_bmi = n_user_demands,
     )
 end

core/test/allocation_test.jl

@@ -8,9 +8,12 @@
     @test ispath(toml_path)
     model = Ribasim.Model(toml_path)
     (; u, p) = model.integrator
-    (; flow_dict) = p.allocation
+    (; input_flow_dict) = p.allocation
 
-    u.flow_allocation_input[flow_dict[(NodeID(:FlowBoundary, 1), NodeID(:Basin, 2))]] = 4.5
+    u.flow_allocation_input[input_flow_dict[(
+        NodeID(:FlowBoundary, 1),
+        NodeID(:Basin, 2),
+    )]] = 4.5
     allocation_model = p.allocation.allocation_models[1]
     Ribasim.allocate!(p, allocation_model, 0.0, u, OptimizationType.allocate)
 
@@ -206,7 +209,7 @@ end
         subnetwork_demands,
         subnetwork_allocateds,
         record_flow,
-        flow_dict,
+        input_flow_dict,
     ) = allocation
 
     # Collecting demands
@@ -236,8 +239,10 @@ end
 
     # Running full allocation algorithm
     (; Δt_allocation) = allocation_models[1]
-    u.flow_allocation_input[flow_dict[(NodeID(:FlowBoundary, 1), NodeID(:Basin, 2))]] =
-        4.5 * Δt_allocation
+    u.flow_allocation_input[input_flow_dict[(
+        NodeID(:FlowBoundary, 1),
+        NodeID(:Basin, 2),
+    )]] = 4.5 * Δt_allocation
     Ribasim.update_allocation!((; p, t, u))
 
     @test subnetwork_allocateds[NodeID(:Basin, 2), NodeID(:Pump, 11)] ≈
@@ -274,13 +279,18 @@ end
     model = Ribasim.Model(toml_path)
     (; p, u, t) = model.integrator
     (; allocation, user_demand, graph, basin) = p
-    (; allocation_models, subnetwork_demands, subnetwork_allocateds, flow_dict) = allocation
+    (; allocation_models, subnetwork_demands, subnetwork_allocateds, input_flow_dict) =
+        allocation
 
     # Set flows of sources in
-    u.flow_allocation_input[flow_dict[(NodeID(:FlowBoundary, 58), NodeID(:Basin, 16))]] =
-        1.0
-    u.flow_allocation_input[flow_dict[(NodeID(:FlowBoundary, 59), NodeID(:Basin, 44))]] =
-        1e-3
+    u.flow_allocation_input[input_flow_dict[(
+        NodeID(:FlowBoundary, 58),
+        NodeID(:Basin, 16),
+    )]] = 1.0
+    u.flow_allocation_input[input_flow_dict[(
+        NodeID(:FlowBoundary, 59),
+        NodeID(:Basin, 44),
+    )]] = 1e-3
 
     # Collecting demands
     for allocation_model in allocation_models[2:end]
@@ -421,7 +431,7 @@ end
     @test constraint_flow_demand_outflow.set.upper == 0.0
 
     optimization_type = OptimizationType.internal_sources
-    for (edge, idx) in allocation.flow_dict
+    for (edge, idx) in allocation.input_flow_dict
         u.flow_allocation_input[idx] = Ribasim.get_flow(graph, edge..., 0)
     end
     Ribasim.set_initial_values!(allocation_model, p, u, t)

core/test/utils_test.jl

@@ -199,7 +199,7 @@ end
     db = SQLite.DB(db_path)
 
     p = Ribasim.Parameters(db, cfg)
-    n_states = sum(Ribasim.get_n_states(db, config))
+    n_states = sum(Ribasim.get_n_states(db, cfg))
     close(db)
 
     jac_prototype = Ribasim.get_jac_prototype(p, n_states)

docs/core/equations.qmd

@@ -37,7 +37,10 @@ Let $V$ be the set of node IDs and let $E$ be the set of ordered tuples $(i,j)$
 We can split the set of nodes into two subsets $V = B \cup N$, where $B$ is the set of basins and $N$ is the set of non-basins.
 The basins have an associated storage state and the non-basins dictate how water flows to or from basins.
 
-$\mathbf{u}(t)$ is given by all the states of the model, which are (currently) the storage of the basins and the integral terms of the PID controllers, the latter being explained in [3 PID controller](equations.qmd#sec-PID).
+$\mathbf{u}(t)$ is given by all the states of the model, which are:
+- the storage of the basins;
+- the integral terms of the PID controllers, further explained in [3 PID controller](equations.qmd#sec-PID);
+- integrated flows $\int_{t_0}^t Q_{i,j}(t')\text{d}t'$ used for calculating mean flows.
 
 Given a single basin with node ID $i \in B$, the equation that dictates the change of its storage over time is given by