Store state node IDs in parameters

core/src/logging.jl

@@ -41,7 +41,7 @@ function log_bottlenecks(model; converged::Bool)
         max_errors = 5
         # Iterate over the errors in descending order
         for i in sortperm(flow_error; rev = true)
-            node_id = Symbol(id_from_state_index(p, u, i))
+            node_id = Symbol(p.node_id[i])
             error = flow_error[i]
             isnan(error) && continue  # NaN are sorted as largest
             # Stop reporting errors if they are too small or too many

core/src/parameter.jl

@@ -1003,6 +1003,8 @@ const ModelGraph = MetaGraph{
     const water_balance_reltol::Float64
     # State at previous saveat
     const u_prev_saveat::C11 = ComponentVector()
+    # Node ID associated with each state
+    const node_id::Vector{NodeID} = NodeID[]
 end
 
 # To opt-out of type checking for ForwardDiff

core/src/read.jl

@@ -1700,6 +1700,20 @@ function Parameters(db::DB, config::Config)::Parameters
 
     subgrid = Subgrid(db, config, basin)
 
+    # aligned with the state
+    node_id = vcat(
+        tabulated_rating_curve.node_id,
+        pump.node_id,
+        outlet.node_id,
+        user_demand.node_id,
+        user_demand.node_id,
+        linear_resistance.node_id,
+        manning_resistance.node_id,
+        basin.node_id,
+        basin.node_id,
+        pid_control.node_id,
+    )
+
     p = Parameters(;
         config.starttime,
         graph,
@@ -1722,6 +1736,7 @@ function Parameters(db::DB, config::Config)::Parameters
         subgrid,
         config.solver.water_balance_abstol,
         config.solver.water_balance_reltol,
+        node_id,
     )
 
     collect_control_mappings!(p)

core/src/util.jl

@@ -867,7 +867,7 @@ end
 function build_state_vector(p::Parameters)
     # It is assumed that the horizontal flow states come first in
     # p.state_inflow_link and p.state_outflow_link
-    return ComponentVector{Float64}(;
+    u = ComponentVector{Float64}(;
         tabulated_rating_curve = zeros(length(p.tabulated_rating_curve.node_id)),
         pump = zeros(length(p.pump.node_id)),
         outlet = zeros(length(p.outlet.node_id)),
@@ -879,6 +879,8 @@ function build_state_vector(p::Parameters)
         infiltration = zeros(length(p.basin.node_id)),
         integral = zeros(length(p.pid_control.node_id)),
     )
+    @assert length(u) == length(p.node_id)
+    return u
 end
 
 function build_flow_to_storage(p::Parameters, u::ComponentVector)::Parameters
@@ -999,32 +1001,6 @@ function set_state_flow_links(p::Parameters, u0::ComponentVector)::Parameters
     return p
 end
 
-function id_from_state_index(
-    p::Parameters,
-    ::ComponentVector{Float64, Vector{Float64}, <:Tuple{<:Axis{NT}}},
-    global_idx::Int,
-)::NodeID where {NT}
-    local_idx = 0
-    component = Symbol()
-    for (comp, range) in pairs(NT)
-        if global_idx in range
-            component = comp
-            local_idx = global_idx - first(range) + 1
-            break
-        end
-    end
-    component_string = String(component)
-    if endswith(component_string, "_inflow") || endswith(component_string, "_outflow")
-        component = :user_demand
-    elseif component == :integral
-        component = :pid_control
-    elseif component in [:infiltration, :evaporation]
-        component = :basin
-    end
-
-    getfield(p, component).node_id[local_idx]
-end
-
 function get_state_index(
     id::NodeID,
     ::ComponentVector{A, B, <:Tuple{<:Axis{NT}}};

core/test/allocation_test.jl

@@ -15,7 +15,7 @@
         4.5
     allocation_model = p.allocation.allocation_models[1]
     (; flow) = allocation_model
-    u = ComponentVector(; storage = zeros(length(p.basin.node_id)))
+    u = ComponentVector()
     t = 0.0
     Ribasim.allocate_demands!(p, allocation_model, t, u)
 
@@ -133,7 +133,7 @@ end
     t = 0.0
 
     # Collecting demands
-    u = ComponentVector(; storage = zeros(length(basin.node_id)))
+    u = ComponentVector()
     for allocation_model in allocation_models[2:end]
         Ribasim.collect_demands!(p, allocation_model, t, u)
     end
@@ -227,7 +227,7 @@ end
     mean_input_flows[4][(NodeID(:FlowBoundary, 59, p), NodeID(:Basin, 44, p))] = 1e-3
 
     # Collecting demands
-    u = ComponentVector(; storage = zeros(length(basin.node_id)))
+    u = ComponentVector()
     for allocation_model in allocation_models[2:end]
         Ribasim.collect_demands!(p, allocation_model, t, u)
     end

core/test/utils_test.jl

@@ -263,6 +263,7 @@ end
     db = SQLite.DB(db_path)
 
     p = Ribasim.Parameters(db, cfg)
+    @test p.node_id == [4, 5, 8, 7, 10, 12, 2, 1, 3, 6, 9, 1, 3, 6, 9]
     close(db)
     t0 = 0.0
     u0 = Ribasim.build_state_vector(p)

docs/concept/allocation.qmd

@@ -241,7 +241,6 @@ The following is an example of an optimization problem for the example shown [he
 using Ribasim
 using Ribasim: NodeID
 using SQLite
-using ComponentArrays: ComponentVector
 
 toml_path = normpath(@__DIR__, "../../generated_testmodels/allocation_example/ribasim.toml")
 model = Ribasim.Model(toml_path)