Merge branch 'main' into add-api

.github/workflows/python_lint.yml

@@ -30,3 +30,6 @@ jobs:
       - name: Run mypy on python/ribasim_api
         run: |
           pixi run mypy-ribasim-api
+      - name: Run mypy on ribasim_qgis
+        run: |
+          pixi run mypy-ribasim-qgis

.teamcity/Ribasim_Windows/buildTypes/Ribasim_Windows_BuildLibribasim.xml

@@ -13,7 +13,8 @@
     <build-runners>
       <runner id="RUNNER_2415" name="Set up pixi" type="simpleRunner">
         <parameters>
-          <param name="script.content" value="pixi run install-without-pre-commit" />
+          <param name="script.content"><![CDATA[pixi --version
+pixi run install-without-pre-commit]]></param>
           <param name="teamcity.build.workingDir" value="ribasim" />
           <param name="teamcity.step.mode" value="default" />
           <param name="use.custom.script" value="true" />

core/src/allocation.jl

@@ -32,8 +32,7 @@ function allocation_graph_used_nodes!(p::Parameters, allocation_network_id::Int)
         use_node = false
         has_fractional_flow_outneighbors =
             get_fractional_flow_connected_basins(node_id, basin, fractional_flow, graph)[3]
-        node_type = graph[node_id].type
-        if node_type in [:user, :basin, :terminal]
+        if node_id.type in [NodeType.User, NodeType.Basin, NodeType.Terminal]
             use_node = true
         elseif has_fractional_flow_outneighbors
             use_node = true
@@ -244,8 +243,7 @@ function process_allocation_graph_edges!(
         # edge are now nodes that have an equivalent in the allocation graph,
         # these do not constrain the composite edge capacity
         for (node_id_1, node_id_2, node_id_3) in IterTools.partition(edge_composite, 3, 1)
-            node_type = graph[node_id_2].type
-            node = getfield(p, node_type)
+            node = getfield(p, graph[node_id_2].type)
 
             # Find flow constraints
             if is_flow_constraining(node)
@@ -281,7 +279,8 @@ function process_allocation_graph_edges!(
     return capacity
 end
 
-const allocation_source_nodetypes = Set{Symbol}([:level_boundary, :flow_boundary])
+const allocation_source_nodetypes =
+    Set{NodeType.T}([NodeType.LevelBoundary, NodeType.FlowBoundary])
 
 """
 Remove allocation user return flow edges that are upstream of the user itself.
@@ -290,7 +289,7 @@ function avoid_using_own_returnflow!(p::Parameters, allocation_network_id::Int):
     (; 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]
+    node_ids_user = [node_id for node_id in node_ids if node_id.type == NodeType.User]
 
     for node_id_user in node_ids_user
         node_id_return_flow = only(outflow_ids_allocation(graph, node_id_user))
@@ -384,7 +383,7 @@ function add_variables_absolute_value!(
     (; main_network_connections) = allocation
     if startswith(config.allocation.objective_type, "linear")
         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]
+        node_ids_user = [node_id for node_id in node_ids if node_id.type == NodeType.User]
 
         # For the main network, connections to subnetworks are treated as users
         if is_main_network(allocation_network_id)
@@ -495,18 +494,18 @@ Add the flow conservation constraints to the allocation problem.
 The constraint indices are user node IDs.
 
 Constraint:
-sum(flows out of node node) <= flows into node + flow from storage and vertical fluxes
+sum(flows out of node node) == flows into node + flow from storage and vertical fluxes
 """
 function add_constraints_flow_conservation!(
     problem::JuMP.Model,
     p::Parameters,
     allocation_network_id::Int,
 )::Nothing
-    (; graph, allocation) = p
+    (; graph) = p
     F = problem[:F]
     node_ids = graph[].node_ids[allocation_network_id]
     node_ids_conservation =
-        [node_id for node_id in node_ids if graph[node_id].type == :basin]
+        [node_id for node_id in node_ids if node_id.type == NodeType.Basin]
     main_network_source_edges = get_main_network_connections(p, allocation_network_id)
     for edge in main_network_source_edges
         push!(node_ids_conservation, edge[2])
@@ -518,7 +517,7 @@ function add_constraints_flow_conservation!(
         sum([
             F[(node_id, outneighbor_id)] for
             outneighbor_id in outflow_ids_allocation(graph, node_id)
-        ]) <= sum([
+        ]) == sum([
             F[(inneighbor_id, node_id)] for
             inneighbor_id in inflow_ids_allocation(graph, node_id)
         ]),
@@ -544,8 +543,8 @@ function add_constraints_user_returnflow!(
 
     node_ids = graph[].node_ids[allocation_network_id]
     node_ids_user_with_returnflow = [
-        node_id for node_id in node_ids if
-        graph[node_id].type == :user && !isempty(outflow_ids_allocation(graph, node_id))
+        node_id for node_id in node_ids if node_id.type == NodeType.User &&
+        !isempty(outflow_ids_allocation(graph, node_id))
     ]
     problem[:return_flow] = JuMP.@constraint(
         problem,
@@ -576,7 +575,7 @@ function add_constraints_absolute_value!(
     objective_type = config.allocation.objective_type
     if startswith(objective_type, "linear")
         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]
+        node_ids_user = [node_id for node_id in node_ids if node_id.type == NodeType.User]
 
         # For the main network, connections to subnetworks are treated as users
         if is_main_network(allocation_network_id)
@@ -655,12 +654,12 @@ function add_constraints_fractional_flow!(
     inflows = Dict{NodeID, JuMP.AffExpr}()
     for node_id in node_ids
         for outflow_id_ in outflow_ids(graph, node_id)
-            if graph[outflow_id_].type == :fractional_flow
+            if outflow_id_.type == NodeType.FractionalFlow
                 # The fractional flow nodes themselves are not represented in
                 # the allocation graph
                 dst_id = outflow_id(graph, outflow_id_)
                 # For now only consider fractional flow nodes which end in a basin
-                if haskey(graph, node_id, dst_id) && graph[dst_id].type == :basin
+                if haskey(graph, node_id, dst_id) && dst_id.type == NodeType.Basin
                     edge = (node_id, dst_id)
                     push!(edges_to_fractional_flow, edge)
                     node_idx = findsorted(fractional_flow.node_id, outflow_id_)
@@ -824,14 +823,11 @@ function set_objective_priority!(
         ex = sum(problem[:F_abs])
     end
 
-    demand_max = 0.0
-
     # Terms for subnetworks as users
     if is_main_network(allocation_network_id)
         for connections_subnetwork in main_network_connections
             for connection in connections_subnetwork
                 d = subnetwork_demands[connection][priority_idx]
-                demand_max = max(demand_max, d)
                 add_user_term!(ex, connection, objective_type, d, allocation_model)
             end
         end
@@ -840,7 +836,7 @@ function set_objective_priority!(
     # Terms for user nodes
     for edge_id in edge_ids
         node_id_user = edge_id[2]
-        if graph[node_id_user].type != :user
+        if node_id_user.type != NodeType.User
             continue
         end
 
@@ -853,25 +849,9 @@ function set_objective_priority!(
             d = get_user_demand(user, node_id_user, priority_idx)
         end
 
-        demand_max = max(demand_max, d)
         add_user_term!(ex, edge_id, objective_type, d, allocation_model)
     end
 
-    # Add flow cost
-    if objective_type == :linear_absolute
-        cost_per_flow = 0.5 / length(F)
-        for flow in F
-            JuMP.add_to_expression!(ex, cost_per_flow * flow)
-        end
-    elseif objective_type == :linear_relative
-        if demand_max > 0.0
-            cost_per_flow = 0.5 / (demand_max * length(F))
-            for flow in F
-                JuMP.add_to_expression!(ex, cost_per_flow * flow)
-            end
-        end
-    end
-
     new_objective = JuMP.@expression(problem, ex)
     JuMP.@objective(problem, Min, new_objective)
     return nothing
@@ -911,7 +891,7 @@ function assign_allocations!(
 
         user_node_id = edge_id[2]
 
-        if graph[user_node_id].type == :user
+        if user_node_id.type == NodeType.User
             allocated = JuMP.value(F[edge_id])
             user_idx = findsorted(user.node_id, user_node_id)
             user.allocated[user_idx][priority_idx] = allocated
@@ -1128,7 +1108,7 @@ function allocate!(
         @debug JuMP.solution_summary(problem)
         if JuMP.termination_status(problem) !== JuMP.OPTIMAL
             (; allocation_network_id) = allocation_model
-            priority = priorities[priority_index]
+            priority = priorities[priority_idx]
             error(
                 "Allocation of subnetwork $allocation_network_id, priority $priority coudn't find optimal solution.",
             )

core/src/callback.jl

@@ -123,29 +123,27 @@ function get_value(
     (; basin, flow_boundary, level_boundary) = p
 
     if variable == "level"
-        hasindex_basin, basin_idx = id_index(basin.node_id, node_id)
-        level_boundary_idx = findsorted(level_boundary.node_id, node_id)
-
-        if hasindex_basin
+        if node_id.type == NodeType.Basin
+            _, basin_idx = id_index(basin.node_id, node_id)
             _, level = get_area_and_level(basin, basin_idx, u[basin_idx])
-        elseif level_boundary_idx !== nothing
+        elseif node_id.type == NodeType.LevelBoundary
+            level_boundary_idx = findsorted(level_boundary.node_id, node_id)
             level = level_boundary.level[level_boundary_idx](t + Δt)
         else
             error(
                 "Level condition node '$node_id' is neither a basin nor a level boundary.",
             )
         end
-
         value = level
 
     elseif variable == "flow_rate"
-        flow_boundary_idx = findsorted(flow_boundary.node_id, node_id)
-
-        if flow_boundary_idx === nothing
+        if node_id.type == NodeType.FlowBoundary
+            flow_boundary_idx = findsorted(flow_boundary.node_id, node_id)
+            value = flow_boundary.flow_rate[flow_boundary_idx](t + Δt)
+        else
             error("Flow condition node $node_id is not a flow boundary.")
         end
 
-        value = flow_boundary.flow_rate[flow_boundary_idx](t + Δt)
     else
         error("Unsupported condition variable $variable.")
     end
@@ -418,7 +416,7 @@ function update_basin(integrator)::Nothing
     )
 
     for row in timeblock
-        hasindex, i = id_index(node_id, NodeID(row.node_id))
+        hasindex, i = id_index(node_id, NodeID(NodeType.Basin, row.node_id))
         @assert hasindex "Table 'Basin / time' contains non-Basin IDs"
         set_table_row!(table, row, i)
     end
@@ -461,7 +459,7 @@ function update_tabulated_rating_curve!(integrator)::Nothing
         id = first(group).node_id
         level = [row.level for row in group]
         flow_rate = [row.flow_rate for row in group]
-        i = searchsortedfirst(node_id, NodeID(id))
+        i = searchsortedfirst(node_id, NodeID(NodeType.TabulatedRatingCurve, id))
         tables[i] = LinearInterpolation(flow_rate, level; extrapolate = true)
     end
     return nothing

core/src/graph.jl

@@ -10,7 +10,7 @@ function create_graph(db::DB, config::Config, chunk_sizes::Vector{Int})::MetaGra
         execute(db, "SELECT fid, type, allocation_network_id FROM Node ORDER BY fid")
     edge_rows = execute(
         db,
-        "SELECT fid, from_node_id, to_node_id, edge_type, allocation_network_id FROM Edge ORDER BY fid",
+        "SELECT fid, from_node_type, from_node_id, to_node_type, to_node_id, edge_type, allocation_network_id FROM Edge ORDER BY fid",
     )
     # Node IDs per subnetwork
     node_ids = Dict{Int, Set{NodeID}}()
@@ -34,7 +34,7 @@ function create_graph(db::DB, config::Config, chunk_sizes::Vector{Int})::MetaGra
         graph_data = nothing,
     )
     for row in node_rows
-        node_id = NodeID(row.fid)
+        node_id = NodeID(row.type, row.fid)
         # Process allocation network ID
         if ismissing(row.allocation_network_id)
             allocation_network_id = 0
@@ -51,15 +51,23 @@ function create_graph(db::DB, config::Config, chunk_sizes::Vector{Int})::MetaGra
             flow_vertical_dict[node_id] = flow_vertical_counter
         end
     end
-    for (; fid, from_node_id, to_node_id, edge_type, allocation_network_id) in edge_rows
+    for (;
+        fid,
+        from_node_type,
+        from_node_id,
+        to_node_type,
+        to_node_id,
+        edge_type,
+        allocation_network_id,
+    ) in edge_rows
         try
             # hasfield does not work
             edge_type = getfield(EdgeType, Symbol(edge_type))
         catch
             error("Invalid edge type $edge_type.")
         end
-        id_src = NodeID(from_node_id)
-        id_dst = NodeID(to_node_id)
+        id_src = NodeID(from_node_type, from_node_id)
+        id_dst = NodeID(to_node_type, to_node_id)
         if ismissing(allocation_network_id)
             allocation_network_id = 0
         end

core/src/parameter.jl

@@ -1,14 +1,35 @@
+# EdgeType.flow and NodeType.FlowBoundary
+@enumx EdgeType flow control none
+@eval @enumx NodeType $(config.nodetypes...)
+
+# Support creating a NodeType enum instance from a symbol or string
+function NodeType.T(s::Symbol)::NodeType.T
+    symbol_map = EnumX.symbol_map(NodeType.T)
+    for (sym, val) in symbol_map
+        sym == s && return NodeType.T(val)
+    end
+    throw(ArgumentError("Invalid value for NodeType: $s"))
+end
+
+NodeType.T(str::AbstractString) = NodeType.T(Symbol(str))
+
 struct NodeID
+    type::NodeType.T
     value::Int
 end
 
+NodeID(type::Symbol, value::Int) = NodeID(NodeType.T(type), value)
+NodeID(type::AbstractString, value::Int) = NodeID(NodeType.T(type), value)
+
 Base.Int(id::NodeID) = id.value
-Base.convert(::Type{NodeID}, value::Int) = NodeID(value)
 Base.convert(::Type{Int}, id::NodeID) = id.value
 Base.broadcastable(id::NodeID) = Ref(id)
-Base.show(io::IO, id::NodeID) = print(io, '#', Int(id))
+Base.show(io::IO, id::NodeID) = print(io, id.type, " #", Int(id))
 
 function Base.isless(id_1::NodeID, id_2::NodeID)::Bool
+    if id_1.type != id_2.type
+        error("Cannot compare NodeIDs of different types")
+    end
     return Int(id_1) < Int(id_2)
 end
 
@@ -64,8 +85,6 @@ struct Allocation
     }
 end
 
-@enumx EdgeType flow control none
-
 """
 Type for storing metadata of nodes in the graph
 type: type of the node
@@ -318,7 +337,7 @@ struct Pump{T} <: AbstractParameterNode
         control_mapping,
         is_pid_controlled,
     ) where {T}
-        if valid_flow_rates(node_id, get_tmp(flow_rate, 0), control_mapping, :Pump)
+        if valid_flow_rates(node_id, get_tmp(flow_rate, 0), control_mapping)
             return new{T}(
                 node_id,
                 active,
@@ -363,7 +382,7 @@ struct Outlet{T} <: AbstractParameterNode
         control_mapping,
         is_pid_controlled,
     ) where {T}
-        if valid_flow_rates(node_id, get_tmp(flow_rate, 0), control_mapping, :Outlet)
+        if valid_flow_rates(node_id, get_tmp(flow_rate, 0), control_mapping)
             return new{T}(
                 node_id,
                 active,