connectivity refactor

core/src/allocation.jl

@@ -3,8 +3,7 @@ Find all nodes in the subnetwork which will be used in the allocation network.
 Some nodes are skipped to optimize allocation optimization.
 """
 function allocation_graph_used_nodes!(p::Parameters, allocation_network_id::Int)::Nothing
-    (; connectivity) = p
-    (; graph) = connectivity
+    (; graph) = p
 
     node_ids = graph[].node_ids[allocation_network_id]
     used_nodes = Set{NodeID}()
@@ -69,8 +68,7 @@ function find_allocation_graph_edges!(
     p::Parameters,
     allocation_network_id::Int,
 )::Tuple{Vector{Vector{NodeID}}, SparseMatrixCSC{Float64, Int}}
-    (; connectivity, user) = p
-    (; graph) = connectivity
+    (; graph) = p
 
     edges_composite = Vector{NodeID}[]
     capacity = spzeros(nv(graph), nv(graph))
@@ -154,8 +152,7 @@ function process_allocation_graph_edges!(
     p::Parameters,
     allocation_network_id::Int,
 )::SparseMatrixCSC{Float64, Int}
-    (; connectivity) = p
-    (; graph) = connectivity
+    (; graph) = p
     node_ids = graph[].node_ids[allocation_network_id]
     edge_ids = graph[].edge_ids[allocation_network_id]
 
@@ -244,8 +241,7 @@ end
 The source nodes must only have one allocation outneighbor and no allocation inneighbors.
 """
 function valid_sources(p::Parameters, allocation_network_id::Int)::Bool
-    (; connectivity) = p
-    (; graph) = connectivity
+    (; graph) = p
 
     edge_ids = graph[].edge_ids[allocation_network_id]
 
@@ -279,8 +275,7 @@ end
 Remove allocation user return flow edges that are upstream of the user itself.
 """
 function avoid_using_own_returnflow!(p::Parameters, allocation_network_id::Int)::Nothing
-    (; connectivity) = p
-    (; graph) = connectivity
+    (; graph) = p
     node_ids = graph[].node_ids[allocation_network_id]
     edge_ids = graph[].edge_ids[allocation_network_id]
     node_ids_user = [node_id for node_id in node_ids if graph[node_id].type == :user]
@@ -334,8 +329,7 @@ function add_variables_flow!(
     p::Parameters,
     allocation_network_id::Int,
 )::Nothing
-    (; connectivity) = p
-    (; graph) = connectivity
+    (; graph) = p
     edge_ids = graph[].edge_ids[allocation_network_id]
     problem[:F] = JuMP.@variable(problem, F[edge_id = edge_ids,] >= 0.0)
     return nothing
@@ -353,8 +347,7 @@ function add_variables_absolute_value!(
     allocation_network_id::Int,
     config::Config,
 )::Nothing
-    (; connectivity) = p
-    (; graph) = connectivity
+    (; graph) = p
     node_ids = graph[].node_ids[allocation_network_id]
     node_ids_user = [node_id for node_id in node_ids if graph[node_id].type == :user]
     if startswith(config.allocation.objective_type, "linear")
@@ -377,8 +370,7 @@ function add_constraints_capacity!(
     p::Parameters,
     allocation_network_id::Int,
 )::Nothing
-    (; connectivity) = p
-    (; graph) = connectivity
+    (; graph) = p
     F = problem[:F]
     edge_ids = graph[].edge_ids[allocation_network_id]
     edge_ids_finite_capacity = Tuple{NodeID, NodeID}[]
@@ -409,8 +401,7 @@ function add_constraints_source!(
     p::Parameters,
     allocation_network_id::Int,
 )::Nothing
-    (; connectivity) = p
-    (; graph) = connectivity
+    (; graph) = p
     edge_ids = graph[].edge_ids[allocation_network_id]
     edge_ids_source = [
         edge_id for edge_id in edge_ids if
@@ -466,8 +457,7 @@ function add_constraints_flow_conservation!(
     p::Parameters,
     allocation_network_id::Int,
 )::Nothing
-    (; connectivity) = p
-    (; graph) = connectivity
+    (; graph) = p
     F = problem[:F]
     node_ids = graph[].node_ids[allocation_network_id]
     node_ids_basin = [node_id for node_id in node_ids if graph[node_id].type == :basin]
@@ -498,8 +488,7 @@ function add_constraints_user_returnflow!(
     p::Parameters,
     allocation_network_id::Int,
 )::Nothing
-    (; connectivity, user) = p
-    (; graph) = connectivity
+    (; graph, user) = p
     F = problem[:F]
 
     node_ids = graph[].node_ids[allocation_network_id]
@@ -531,8 +520,7 @@ function add_constraints_absolute_value!(
     allocation_network_id::Int,
     config::Config,
 )::Nothing
-    (; connectivity) = p
-    (; graph) = connectivity
+    (; graph) = p
     node_ids = graph[].node_ids[allocation_network_id]
 
     objective_type = config.allocation.objective_type
@@ -658,8 +646,7 @@ function set_objective_priority!(
     priority_idx::Int,
 )::Nothing
     (; objective_type, problem, allocation_network_id) = allocation_model
-    (; connectivity, user) = p
-    (; graph) = connectivity
+    (; graph, user) = p
     (; demand, node_id) = user
     edge_ids = graph[].edge_ids[allocation_network_id]
 
@@ -731,12 +718,10 @@ function assign_allocations!(
     priority_idx::Int,
 )::Nothing
     (; problem, allocation_network_id) = allocation_model
-    (; connectivity, user) = p
-    (; graph, flow) = connectivity
+    (; graph, user) = p
     (; record) = user
     edge_ids = graph[].edge_ids[allocation_network_id]
     F = problem[:F]
-    flow = get_tmp(flow, 0)
     for edge_id in edge_ids
         user_node_id = edge_id[2]
         if graph[user_node_id].type != :user
@@ -755,7 +740,10 @@ function assign_allocations!(
         push!(record.allocated, allocated)
         # TODO: This is now the last abstraction before the allocation update,
         # should be the average abstraction since the last allocation solve
-        push!(record.abstracted, flow[inflow_id(graph, user_node_id), user_node_id])
+        push!(
+            record.abstracted,
+            get_flow(graph, inflow_id(graph, user_node_id), user_node_id, 0),
+        )
     end
     return nothing
 end
@@ -769,22 +757,22 @@ function adjust_source_flows!(
     priority_idx::Int,
 )::Nothing
     (; problem) = allocation_model
-    (; connectivity) = p
-    (; graph) = connectivity
+    (; graph) = p
     (; allocation_network_id) = allocation_model
     edge_ids = graph[].edge_ids[allocation_network_id]
     source_constraints = problem[:source]
     F = problem[:F]
 
     # It is assumed that the allocation procedure does not have to be differentiated.
-    flow = get_tmp(p.connectivity.flow, 0)
-
     for edge_id in edge_ids
         # If it is a source edge.
         if graph[edge_id...].allocation_network_id_source == allocation_network_id
             if priority_idx == 1
                 # Reset the source to the current flow.
-                JuMP.set_normalized_rhs(source_constraints[edge_id], flow[edge_id...])
+                JuMP.set_normalized_rhs(
+                    source_constraints[edge_id],
+                    get_flow(graph, edge_id..., 0),
+                )
             else
                 # Subtract the allocated flow from the source.
                 JuMP.set_normalized_rhs(
@@ -809,8 +797,7 @@ function adjust_edge_capacities!(
     p::Parameters,
     priority_idx::Int,
 )::Nothing
-    (; connectivity) = p
-    (; graph) = connectivity
+    (; graph) = p
     (; problem, capacity, allocation_network_id) = allocation_model
     edge_ids = graph[].edge_ids[allocation_network_id]
     constraints_capacity = problem[:capacity]

core/src/bmi.jl

@@ -40,19 +40,19 @@ function BMI.initialize(T::Type{Model}, config::Config)::Model
             error("Invalid discrete control state definition(s).")
         end
 
-        (; pid_control, connectivity, basin, pump, outlet, fractional_flow) = parameters
+        (; pid_control, basin, pump, graph, fractional_flow) = parameters
         if !valid_pid_connectivity(
             pid_control.node_id,
             pid_control.listen_node_id,
-            connectivity.graph,
+            graph,
             basin.node_id,
             pump.node_id,
         )
             error("Invalid PidControl connectivity.")
         end
 
         if !valid_fractional_flow(
-            connectivity.graph,
+            graph,
             fractional_flow.node_id,
             fractional_flow.control_mapping,
         )
@@ -414,7 +414,7 @@ function discrete_control_affect!(
     upcrossing::Union{Bool, Missing},
 )::Bool
     p = integrator.p
-    (; discrete_control, connectivity) = p
+    (; discrete_control, graph) = p
 
     # Get the discrete_control node that listens to this condition
     discrete_control_node_id = discrete_control.node_id[condition_idx]
@@ -474,11 +474,8 @@ function discrete_control_affect!(
         push!(record.control_state, control_state_new)
 
         # Loop over nodes which are under control of this control node
-        for target_node_id in outneighbor_labels_type(
-            connectivity.graph,
-            discrete_control_node_id,
-            EdgeType.control,
-        )
+        for target_node_id in
+            outneighbor_labels_type(graph, discrete_control_node_id, EdgeType.control)
             set_control_params!(p, target_node_id, control_state_new)
         end
 
@@ -489,7 +486,7 @@ function discrete_control_affect!(
 end
 
 function set_control_params!(p::Parameters, node_id::NodeID, control_state::String)
-    node = getfield(p, p.connectivity.graph[node_id].type)
+    node = getfield(p, p.graph[node_id].type)
     idx = searchsortedfirst(node.node_id, node_id)
     new_state = node.control_mapping[(node_id, control_state)]
 
@@ -503,7 +500,10 @@ end
 
 "Copy the current flow to the SavedValues"
 function save_flow(u, t, integrator)
-    copy(nonzeros(get_tmp(integrator.p.connectivity.flow, u)))
+    [
+        get_tmp(integrator.p.graph[].flow_vertical, 0.0)...,
+        get_tmp(integrator.p.graph[].flow, 0.0)...,
+    ]
 end
 
 function update_subgrid_level!(integrator)::Nothing
@@ -548,7 +548,7 @@ end
 "Solve the allocation problem for all users and assign allocated abstractions to user nodes."
 function update_allocation!(integrator)::Nothing
     (; p, t) = integrator
-    for allocation_model in integrator.p.connectivity.allocation_models
+    for allocation_model in integrator.p.allocation_models
         allocate!(p, allocation_model, t)
     end
 end

core/src/create.jl

@@ -196,49 +196,8 @@ end
 const nonconservative_nodetypes =
     Set{String}(["Basin", "LevelBoundary", "FlowBoundary", "Terminal", "User"])
 
-function Connectivity(db::DB, config::Config, chunk_size::Int)::Connectivity
-    if !valid_edge_types(db)
-        error("Invalid edge types found.")
-    end
-
-    graph = create_graph(db)
-
-    # Build the sparsity structure of the flow matrix
-    flow = spzeros(nv(graph), nv(graph))
-    for e in edges(graph)
-        id_src = label_for(graph, e.src)
-        id_dst = label_for(graph, e.dst)
-        edge_metadata = graph[id_src, id_dst]
-        if edge_metadata.type == EdgeType.flow
-            flow[id_src, id_dst] = 1.0
-        end
-    end
-
-    # Add a self-loop, i.e. an entry on the diagonal, for all non-conservative node types.
-    # This is used to store the gain (positive) or loss (negative) for the water balance.
-    # Note that this only affects the sparsity structure.
-    # We want to do it here to avoid changing that during the simulation and keeping it predictable,
-    # e.g. if we wouldn't do this, inactive nodes can appear if control turns them on during runtime.
-    for (i, nodetype) in enumerate(get_nodetypes(db))
-        if nodetype in nonconservative_nodetypes
-            flow[i, i] = 1.0
-        end
-    end
-    flow .= 0.0
-
-    if config.solver.autodiff
-        # FixedSizeDiffCache performs better for sparse matrix
-        flow = FixedSizeDiffCache(flow, chunk_size)
-    end
-
-    allocation_models = AllocationModel[]
-
-    return Connectivity(graph, flow, allocation_models)
-end
-
 function generate_allocation_models!(p::Parameters, config::Config)::Nothing
-    (; connectivity) = p
-    (; allocation_models, graph) = connectivity
+    (; graph, allocation_models) = p
 
     for allocation_network_id in keys(graph[].node_ids)
         push!(
@@ -771,8 +730,6 @@ function User(db::DB, config::Config)::User
         abstracted = Vector{Float64}(),
     )
 
-    allocation_optimized = BitVector(zeros(UInt8, length(node_ids)))
-
     return User(
         NodeID.(node_ids),
         active,
@@ -821,8 +778,8 @@ end
 function Parameters(db::DB, config::Config)::Parameters
     n_states = length(get_ids(db, "Basin")) + length(get_ids(db, "PidControl"))
     chunk_size = pickchunksize(n_states)
-
-    connectivity = Connectivity(db, config, chunk_size)
+    graph = create_graph(db, config, chunk_size)
+    allocation_models = Vector{AllocationModel}()
 
     linear_resistance = LinearResistance(db, config)
     manning_resistance = ManningResistance(db, config)
@@ -842,8 +799,7 @@ function Parameters(db::DB, config::Config)::Parameters
 
     # Set is_pid_controlled to true for those pumps and outlets that are PID controlled
     for id in pid_control.node_id
-        id_controlled =
-            only(outneighbor_labels_type(connectivity.graph, id, EdgeType.control))
+        id_controlled = only(outneighbor_labels_type(graph, id, EdgeType.control))
         pump_idx = findsorted(pump.node_id, id_controlled)
         if pump_idx === nothing
             outlet_idx = findsorted(outlet.node_id, id_controlled)
@@ -855,7 +811,8 @@ function Parameters(db::DB, config::Config)::Parameters
 
     p = Parameters(
         config.starttime,
-        connectivity,
+        graph,
+        allocation_models,
         basin,
         linear_resistance,
         manning_resistance,
@@ -872,7 +829,6 @@ function Parameters(db::DB, config::Config)::Parameters
         Dict{Int, Symbol}(),
         subgrid_level,
     )
-
     # Allocation data structures
     if config.allocation.use_allocation
         generate_allocation_models!(p, config)