diff --git a/core/src/allocation.jl b/core/src/allocation.jl index 82018d55e..090eca2cf 100644 --- a/core/src/allocation.jl +++ b/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) @@ -506,7 +505,7 @@ function add_constraints_flow_conservation!( 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]) @@ -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_) @@ -837,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 @@ -892,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 diff --git a/core/src/validation.jl b/core/src/validation.jl index 0ce8fc846..c84322fa1 100644 --- a/core/src/validation.jl +++ b/core/src/validation.jl @@ -551,10 +551,10 @@ function valid_discrete_control(p::Parameters, config::Config)::Bool end for (Δt, var, node_id) in zip(look_ahead, variable, listen_node_id) if !iszero(Δt) - node_type = graph[node_id].type + node_type = node_id.type # TODO: If more transient listen variables must be supported, this validation must be more specific # (e.g. for some node some variables are transient, some not). - if node_type ∉ [:flow_boundary, :level_boundary] + if node_type ∉ [NodeType.FlowBoundary, NodeType.LevelBoundary] errors = true @error "Look ahead supplied for non-timeseries listen variable '$var' from listen node $node_id." else @@ -562,8 +562,8 @@ function valid_discrete_control(p::Parameters, config::Config)::Bool errors = true @error "Negative look ahead supplied for listen variable '$var' from listen node $node_id." else - node = getfield(p, node_type) - idx = if node_type == :Basin + node = getfield(p, graph[node_id].type) + idx = if node_type == NodeType.Basin id_index(node.node_id, node_id) else searchsortedfirst(node.node_id, node_id) @@ -593,7 +593,7 @@ function valid_sources(p::Parameters, allocation_network_id::Int)::Bool for edge in edge_ids (id_source, id_dst) = edge if graph[id_source, id_dst].allocation_network_id_source == allocation_network_id - from_source_node = graph[id_source].type in allocation_source_nodetypes + from_source_node = id_source.type in allocation_source_nodetypes if is_main_network(allocation_network_id) if !from_source_node