Modify main_network_with_subnetworks test model and associated tests

core/src/allocation_init.jl

@@ -108,11 +108,15 @@ This loop finds allocation network edges in several ways:
 function find_allocation_graph_edges!(
     p::Parameters,
     allocation_network_id::Int32,
-)::Tuple{Vector{Vector{NodeID}}, SparseMatrixCSC{Float64, Int}}
+)::Tuple{
+    Vector{Vector{NodeID}},
+    JuMP.Containers.SparseAxisArray{Float64, 2, Tuple{NodeID, NodeID}},
+}
     (; graph) = p
 
     edges_composite = Vector{NodeID}[]
-    capacity = spzeros(nv(graph), nv(graph))
+    dict = Dict{Tuple{NodeID, NodeID}, Float64}()
+    capacity = JuMP.Containers.SparseAxisArray(dict)
 
     node_ids = graph[].node_ids[allocation_network_id]
     edge_ids = Set{Tuple{NodeID, NodeID}}()
@@ -199,11 +203,11 @@ For the composite allocation network edges:
 - Find out their allowed flow direction(s)
 """
 function process_allocation_graph_edges!(
-    capacity::SparseMatrixCSC{Float64, Int},
+    capacity::JuMP.Containers.SparseAxisArray{Float64, 2, Tuple{NodeID, NodeID}},
     edges_composite::Vector{Vector{NodeID}},
     p::Parameters,
     allocation_network_id::Int32,
-)::SparseMatrixCSC{Float64, Int}
+)::Nothing
     (; graph) = p
     node_ids = graph[].node_ids[allocation_network_id]
     edge_ids = graph[].edge_ids[allocation_network_id]
@@ -289,7 +293,7 @@ function process_allocation_graph_edges!(
             push!(edge_ids, (node_id_2, node_id_1))
         end
     end
-    return capacity
+    return nothing
 end
 
 const allocation_source_nodetypes =
@@ -320,7 +324,7 @@ Build the graph used for the allocation problem.
 function allocation_graph(
     p::Parameters,
     allocation_network_id::Int32,
-)::SparseMatrixCSC{Float64, Int}
+)::JuMP.Containers.SparseAxisArray{Float64, 2, Tuple{NodeID, NodeID}}
     # Find out which nodes in the subnetwork are used in the allocation network
     allocation_graph_used_nodes!(p, allocation_network_id)
 
@@ -462,7 +466,7 @@ flow over edge <= edge capacity
 """
 function add_constraints_capacity!(
     problem::JuMP.Model,
-    capacity::SparseMatrixCSC{Float64, Int},
+    capacity::JuMP.Containers.SparseAxisArray{Float64, 2, Tuple{NodeID, NodeID}},
     p::Parameters,
     allocation_network_id::Int32,
 )::Nothing
@@ -472,14 +476,14 @@ function add_constraints_capacity!(
     edge_ids = graph[].edge_ids[allocation_network_id]
     edge_ids_finite_capacity = Tuple{NodeID, NodeID}[]
     for edge in edge_ids
-        if !isinf(capacity[edge...]) && edge ∉ main_network_source_edges
+        if !isinf(get(capacity, edge...)) && edge ∉ main_network_source_edges
             push!(edge_ids_finite_capacity, edge)
         end
     end
     problem[:capacity] = JuMP.@constraint(
         problem,
         [edge = edge_ids_finite_capacity],
-        F[edge] <= capacity[edge...],
+        F[edge] <= get(capacity, edge...),
         base_name = "capacity"
     )
     return nothing
@@ -911,7 +915,7 @@ Construct the allocation problem for the current subnetwork as a JuMP.jl model.
 """
 function allocation_problem(
     p::Parameters,
-    capacity::SparseMatrixCSC{Float64, Int},
+    capacity::JuMP.Containers.SparseAxisArray{Float64, 2, Tuple{NodeID, NodeID}},
     allocation_network_id::Int32,
 )::JuMP.Model
     optimizer = JuMP.optimizer_with_attributes(HiGHS.Optimizer, "log_to_console" => false)

core/src/allocation_optim.jl

@@ -282,7 +282,7 @@ function adjust_capacities_edge!(
     main_network_source_edges = get_main_network_connections(p, allocation_network_id)
 
     for edge_id in edge_ids
-        c = capacity[edge_id...]
+        c = get(capacity, edge_id...)
 
         # These edges have no capacity constraints:
         # - With infinite capacity

core/src/callback.jl

@@ -178,7 +178,10 @@ function get_value(
 
     if variable == "level"
         if node_id.type == NodeType.Basin
-            _, basin_idx = id_index(basin.node_id, node_id)
+            has_index, basin_idx = id_index(basin.node_id, node_id)
+            if !has_index
+                error("Discrete control listen node $node_id does not exist.")
+            end
             _, level = get_area_and_level(basin, basin_idx, u[basin_idx])
         elseif node_id.type == NodeType.LevelBoundary
             level_boundary_idx = findsorted(level_boundary.node_id, node_id)

core/src/model.jl

@@ -85,13 +85,6 @@ function Model(config::Config)::Model
         state = load_structvector(db, config, BasinStateV1)
         n = length(get_ids(db, "Basin"))
 
-        sql = "SELECT node_id FROM Node ORDER BY node_id"
-        node_id = only(execute(columntable, db, sql))
-        if !allunique(node_id)
-            error(
-                "Node IDs need to be globally unique until https://github.com/Deltares/Ribasim/issues/1262 is fixed.",
-            )
-        end
     finally
         # always close the database, also in case of an error
         close(db)

core/src/parameter.jl

@@ -50,7 +50,7 @@ problem: The JuMP.jl model for solving the allocation problem
 """
 struct AllocationModel
     allocation_network_id::Int32
-    capacity::SparseMatrixCSC{Float64, Int}
+    capacity::JuMP.Containers.SparseAxisArray{Float64, 2, Tuple{NodeID, NodeID}}
     problem::JuMP.Model
     Δt_allocation::Float64
 end

core/src/util.jl

@@ -707,3 +707,11 @@ function Base.get(
         nothing
     end
 end
+
+function Base.get(
+    array::JuMP.Containers.SparseAxisArray,
+    node_id_src::NodeID,
+    node_id_dst::NodeID,
+)::Float64
+    get(array.data, (node_id_src, node_id_dst), 0.0)
+end

core/test/allocation_test.jl

@@ -167,34 +167,34 @@ end
 
     # Connections from main network to subnetworks
     @test isempty(main_network_connections[1])
-    @test only(main_network_connections[2]) == (NodeID(:Basin, 2), NodeID(:Pump, 11))
-    @test only(main_network_connections[3]) == (NodeID(:Basin, 6), NodeID(:Pump, 24))
-    @test only(main_network_connections[4]) == (NodeID(:Basin, 10), NodeID(:Pump, 38))
+    @test only(main_network_connections[2]) == (NodeID(:Basin, 100), NodeID(:Pump, 100))
+    @test only(main_network_connections[3]) == (NodeID(:Basin, 300), NodeID(:Pump, 300))
+    @test only(main_network_connections[4]) == (NodeID(:Basin, 500), NodeID(:Pump, 400))
 
     # main-sub connections are part of main network allocation network
     allocation_edges_main_network = graph[].edge_ids[1]
     @test [
-        (NodeID(:Basin, 2), NodeID(:Pump, 11)),
-        (NodeID(:Basin, 6), NodeID(:Pump, 24)),
-        (NodeID(:Basin, 10), NodeID(:Pump, 38)),
+        (NodeID(:Basin, 100), NodeID(:Pump, 100)),
+        (NodeID(:Basin, 300), NodeID(:Pump, 300)),
+        (NodeID(:Basin, 500), NodeID(:Pump, 400)),
     ] ⊆ allocation_edges_main_network
 
     # Subnetworks interpreted as user_demands require variables and constraints to
     # support absolute value expressions in the objective function
     allocation_model_main_network = Ribasim.get_allocation_model(p, Int32(1))
     problem = allocation_model_main_network.problem
-    @test problem[:F_abs_user_demand].axes[1] == NodeID.(:Pump, [11, 24, 38])
-    @test problem[:abs_positive_user_demand].axes[1] == NodeID.(:Pump, [11, 24, 38])
-    @test problem[:abs_negative_user_demand].axes[1] == NodeID.(:Pump, [11, 24, 38])
+    @test problem[:F_abs_user_demand].axes[1] == NodeID.(:Pump, [100, 300, 400])
+    @test problem[:abs_positive_user_demand].axes[1] == NodeID.(:Pump, [100, 300, 400])
+    @test problem[:abs_negative_user_demand].axes[1] == NodeID.(:Pump, [100, 300, 400])
 
     # In each subnetwork, the connection from the main network to the subnetwork is
     # interpreted as a source
     @test Ribasim.get_allocation_model(p, Int32(3)).problem[:source].axes[1] ==
-          [(NodeID(:Basin, 2), NodeID(:Pump, 11))]
+          [(NodeID(:Basin, 100), NodeID(:Pump, 100))]
     @test Ribasim.get_allocation_model(p, Int32(5)).problem[:source].axes[1] ==
-          [(NodeID(:Basin, 6), NodeID(:Pump, 24))]
+          [(NodeID(:Basin, 300), NodeID(:Pump, 300))]
     @test Ribasim.get_allocation_model(p, Int32(7)).problem[:source].axes[1] ==
-          [(NodeID(:Basin, 10), NodeID(:Pump, 38))]
+          [(NodeID(:Basin, 500), NodeID(:Pump, 400))]
 end
 
 @testitem "allocation with main network optimization problem" begin
@@ -224,9 +224,9 @@ end
         Ribasim.allocate!(p, allocation_model, t, u; collect_demands = true)
     end
 
-    @test subnetwork_demands[(NodeID(:Basin, 2), NodeID(:Pump, 11))] ≈ [4.0, 4.0, 0.0]
-    @test subnetwork_demands[(NodeID(:Basin, 6), NodeID(:Pump, 24))] ≈ [0.004, 0.0, 0.0]
-    @test subnetwork_demands[(NodeID(:Basin, 10), NodeID(:Pump, 38))] ≈
+    @test subnetwork_demands[(NodeID(:Basin, 100), NodeID(:Pump, 100))] ≈ [4.0, 4.0, 0.0]
+    @test subnetwork_demands[(NodeID(:Basin, 300), NodeID(:Pump, 300))] ≈ [0.004, 0.0, 0.0]
+    @test subnetwork_demands[(NodeID(:Basin, 500), NodeID(:Pump, 400))] ≈
           [0.001, 0.002, 0.0011]
 
     # Solving for the main network, containing subnetworks as UserDemands
@@ -238,20 +238,20 @@ end
     objective = JuMP.objective_function(problem)
     objective_variables = keys(objective.terms)
     F_abs_user_demand = problem[:F_abs_user_demand]
-    @test F_abs_user_demand[NodeID(:Pump, 11)] ∈ objective_variables
-    @test F_abs_user_demand[NodeID(:Pump, 24)] ∈ objective_variables
-    @test F_abs_user_demand[NodeID(:Pump, 38)] ∈ objective_variables
+    @test F_abs_user_demand[NodeID(:Pump, 100)] ∈ objective_variables
+    @test F_abs_user_demand[NodeID(:Pump, 300)] ∈ objective_variables
+    @test F_abs_user_demand[NodeID(:Pump, 400)] ∈ objective_variables
 
     # Running full allocation algorithm
-    Ribasim.set_flow!(graph, NodeID(:FlowBoundary, 1), NodeID(:Basin, 2), 4.5)
+    Ribasim.set_flow!(graph, NodeID(:FlowBoundary, 100), NodeID(:Basin, 100), 4.5)
     u = ComponentVector(; storage = zeros(length(p.basin.node_id)))
     Ribasim.update_allocation!((; p, t, u))
 
-    @test subnetwork_allocateds[NodeID(:Basin, 2), NodeID(:Pump, 11)] ≈
+    @test subnetwork_allocateds[NodeID(:Basin, 100), NodeID(:Pump, 100)] ≈
           [4.0, 0.49500000, 0.0]
-    @test subnetwork_allocateds[NodeID(:Basin, 6), NodeID(:Pump, 24)] ≈
+    @test subnetwork_allocateds[NodeID(:Basin, 300), NodeID(:Pump, 300)] ≈
           [0.00399999999, 0.0, 0.0]
-    @test subnetwork_allocateds[NodeID(:Basin, 10), NodeID(:Pump, 38)] ≈ [0.001, 0.0, 0.0]
+    @test subnetwork_allocateds[NodeID(:Basin, 500), NodeID(:Pump, 400)] ≈ [0.001, 0.0, 0.0]
 
     @test user_demand.allocated[2] ≈ [4.0, 0.0, 0.0]
     @test user_demand.allocated[7] ≈ [0.001, 0.0, 0.0]

python/ribasim/ribasim/model.py

@@ -3,7 +3,6 @@
 from pathlib import Path
 from typing import Any
 
-import geopandas as gpd
 import numpy as np
 import pandas as pd
 import tomli
@@ -146,16 +145,6 @@ def _save(self, directory: DirectoryPath, input_dir: DirectoryPath):
         for sub in self._nodes():
             sub._save(directory, input_dir)
 
-        # Temporarily require unique node_id for #1262
-        # and copy them to the fid for #1306.
-        df = gpd.read_file(db_path, layer="Node")
-        if not df["node_id"].is_unique:
-            raise ValueError("node_id must be unique")
-        df.set_index("node_id", drop=False, inplace=True)
-        df.sort_index(inplace=True)
-        df.index.name = "fid"
-        df.to_file(db_path, layer="Node", driver="GPKG", index=True)
-
     def node_table(self) -> NodeTable:
         """Compute the full NodeTable from all node types."""
         df_chunks = [node.node.df for node in self._nodes()]