Allocation

1 Overview of allocation implementation

In this document, the allocation workflow is explained. Below is an overview of it.

flowchart TD
    subgraph update_allocation
        direction TB
        G[update mean flows]-->E[collect demand]
        E-->F[allocate]
    end
    style update_allocation fill:#9ff
    C(Begin)-->A[Initialization allocation]
    A--> update_allocation --> H[\end of time?\]
    H--> |No| update_allocation
    H --> |Yes| D(End)

If allocation is used in a model, Allocation structs are created. The allocation struct stores the data that is needed for the calculations and stores also the results of the calculation. In allocation, optimization is an essential part. JuMP.jl is used to model and solve the optimization problems that are defined by allocation. The AllocationModel struct is used for constructing the JuMP model. When an instance of AllocationModel is created, a JuMP optimization model is defined and initialized in the instance. More details on how allocation interacts with JuMP.jl is explained here.

After initialization, as the simulation starts, the allocation problem is solved and updated after every allocation timestep (which is specified in the TOML). With every allocation timestep a new optimization problem is formulated and solved, using the latest available (simulation) model conditions and forcing and demand predictions.

The update of allocation (update_allocation) is repeating and spread into three parts:

• Updating the mean flows. The mean flow data is only used only for output, not used by any internal functions.
• “Collect demand”. This step initialize and solve the optimization problems that collects the demand from the subnetworks.
• “Allocate”. This step solves the optimization problems that allocates the demand. For the main network this step allocates to the subnetworks and demand nodes that are in the main network. For the subnetwork this step allocates to the demand nodes.

The steps “collect demand” and “allocate” correspond to the function collect_demand and allocate_demand in the code.

The iteration stops when it reaches the end time step.

1.1 The Allocation struct

The Allocation struct stores necessary data and calculation results.

field	type	description
subnetwork_ids	Vector{Int32}	The unique sorted allocation network IDs
allocation_models	AllocationModel	The allocation models for the main network and subnetworks corresponding to subnetwork_ids
main_network_connections	Vector{Vector{Tuple{NodeID, NodeID}}}	(from_id, to_id) from the main network to the subnetwork per subnetwork
priorities	Vector{Int32}	All used priority values.
subnetwork_demands	Dict{Tuple{NodeID, NodeID}, Vector{Float64}}	The demand of an edge from the main network to a subnetwork
subnetwork_allocateds	Dict{Tuple{NodeID, NodeID}, Vector{Float64}}	The allocated flow of an edge from the main network to a subnetwork
mean_input_flows	Dict{Tuple{NodeID, NodeID}, Float64}	Flows averaged over Δt_allocation over edges that are allocation sources
mean_realized_flows	Dict{Tuple{NodeID, NodeID}, Float64}	Flows averaged over Δt_allocation over edges that realize a demand
record_demand		A record of demands and allocated flows for nodes that have these
record_flow		A record of all flows computed by allocation optimization, eventually saved to output file

1.2 The AllocationModel struct

The AllocationModel struct has all the data that is needed for the JuMP optimization problem.

field	type	description
subnetwork_id	Int32	The ID of this allocation network
capacity	JuMP.Containers.SparseAxisArray	The capacity per edge of the allocation network, as constrained by nodes that have a max_flow_rate
problem	JuMP.Model	The JuMP.jl model for solving the allocation problem
Δt_allocation	Float64	The time interval between consecutive allocation solves

1.3 JuMP problem interaction

When working with optimization problems using JuMP, there are three fundamental components that need to be defined:

• Optimization variables: These are the variables that are optimized in the allocation problem formulation. They are defined using the @variable macro. For example, to specify the flow rates in all the edges in the allocation network as variables:

problem[:F] = JuMP.@variable(problem, F[edge = edges] >= 0.0)

More details about setting up variables in allocation can be found in the section below.

• Constraints: These are the constraints that the optimization variables must satisfy. They are defined using the @constraint macro. The definition of the edge capacity constraints is shown in section below. add_constraints_... functions are used to add constraints to the optimization problem. The initial value of the constraints is set in the function set_initial_values_.... During the iteration, the constraints are updated based on the current state of the allocation network. When looping over priorities, the constraints are updated by the function adjust_....

• Objective function: This is the function that sets the objective of the optimization problem. It is defined using the @objective macro.

The functions JuMP.normalized_rhs and JuMP.set_normalized_rhs are used to read and write the constant right hand side of constraints.

For example, to update the capacity of one of the edges, JuMP.normalized_rhs moves all the constants to the right-hand sides and all variables to the left-hand side and JuMP.set_normalized_rhs sets the new right-hand-side value.

JuMP.set_normalized_rhs(
    constraints_capacity[edge_id],
    JuMP.normalized_rhs(constraints_capacity[edge_id]) - JuMP.value(F[edge_id]),
)

Some JuMP data structures are used to store intermediate or result data. For more information, see JuMP API.

2 Initialization

Initialization of the allocation data structures happens in allocation_init.jl. Below the steps of allocation problem initialization are explained.

For each subnetwork, an allocation problem is formulated, which is stored in the allocation_models field mentioned above.

2.1 Data processing

2.1.1 Deriving edge capacities

Edge capacities are important constraints in the optimization problem. They set the limit for the flows between the nodes. Therefore, the capacities of all the flow edges in the subnetworks are obtained. The capacity of an edge is given by the smallest max_flow_rate of the nodes connected to the edges if these nodes have such a value. The capacities are stored in a SparseArray object from JuMP.jl called capacities, indexed by a tuple of node IDs.

The function get_capacity obtains the capacities of the edges within a subnetwork given a subnetwork ID and the Ribasim model parameters p, if the sources of the subnetwork are valid (checked in function valid_sources).

2.1.2 Handling the connection between the main network and subnetworks

The function find_subnetwork_connetions finds the edges that connected the main network to a subnetwork. subnetwork_demands and subnetwork_allocateds will be created, which stores demands and allocated values for subnetworks as a whole. main_network_connections is a vector of edges that connect a subnetwork with the main network.

2.2 The optimization problem

2.2.1 Setting up the optimization variables

There are three types of variables in the optimization problems:

• flows between the edges in the allocation model
• flows in and out of a basin with a level demand
• flows in and out of nodes that have a buffer, which are nodes that have a flow demand

The function add_variables_flow is used to add the variable of flows between the edges. The variables are obtained from the capacity array. And variables named by F($startnode, $endnode) are created.

edges = keys(capacity.data)
problem[:F] = JuMP.@variable(problem, F[edge = edges] >= 0.0)

In the function add_variables_basin, variables that represent flows of those basins that are connected with level demand are defined. Part of the function is shown in the code block below. A variable is named F_basin_in if the corresponding basin is supplied by a level demand and F_basin_out if consumed by a level demand.

# Get the node IDs from the subnetwork for basins that have a level demand
node_ids_basin = [
    node_id for
    node_id in graph[].node_ids[subnetwork_id] if graph[node_id].type == :basin &&
    has_external_demand(graph, node_id, :level_demand)[1]
]
problem[:F_basin_in] =
    JuMP.@variable(problem, F_basin_in[node_id = node_ids_basin,] >= 0.0)
problem[:F_basin_out] =
    JuMP.@variable(problem, F_basin_out[node_id = node_ids_basin,] >= 0.0)

The last set of optimization variables is the flow edges in and out of the buffer of nodes with a flow demand. It is defined in a similar way to the second set of variables.

2.2.2 Setting up initial optimization constraints

All the variables are greater and equal to 0. This is set when the variables are added to the optimization problem.

Other constraints are capacity, source_user, source, flow_conservation, fractional_flow, basin_outflow, flow_buffer_outflow and flow_demand_outflow.

For each set of constraints, a function named add_constrains_[constraints name] is created.

Take add_constraints_user_source as an example, the nodes that are relevant for the constraints are added to the optimization problem by calling JuMP.@constraint.

node_ids_user = [node_id for node_id in node_ids if node_id.type == NodeType.UserDemand]

problem[:source_user] = JuMP.@constraint(
    problem,
    [node_id = node_ids_user],
    F[(node_id, outflow_id(graph, node_id))] <= 0.0,
    base_name = "source_user"
)

3 Optimization

Initialization of the data structure is in allocation_init.jl, and updating, running and reading the results is in allocation_optim.jl.

3.1 Preparing the optimization problem

3.1.1 Setting up the objective function

The optimization objective is the sum of three quadratic error terms. The quadratic terms are defined with the add_objective_term function.

Function set_objective_priority sets the objective function based on the main network for a given priority with the following steps:

• First, it treats the subnetworks as user demand nodes and adds the quadratic terms of the main network.
• Then it loops over all the edges in allocation.
• Based on the type of the node that the edge is pointing to (user demand or flow demand), it adds the corresponding quadratic terms.
• Finally, it does the same to the edges that start from a level demand node.

3.1.2 Setting the constraints and capacities

In the function set_initial_values, the following capacities and demands are initialized:

• Source capacities come from the physical layer
• Edge capacities derived from the maximum capacities between the connected nodes
• Basin capacities come from the disk of water above the max level set by a level demand node
• Buffer capacities start at 0
• User demands fractional return flow starts at 0
• Demands either come from the Ribasim model or are set via the BMI

As shown below, these functions set the capacities to the corresponding initial values.

set_initial_capacities_source!(allocation_model, p)
set_initial_capacities_edge!(allocation_model, p)
set_initial_capacities_basin!(allocation_model, p, u, t)
set_initial_capacities_buffer!(allocation_model)
set_initial_capacities_returnflow!(allocation_model)

set_initial_demands_user!(allocation_model, p, t)
set_initial_demands_level!(allocation_model, u, p, t)
set_initial_demands_flow!(allocation_model, p, t)

These capacities determine the constraints of the optimization problem. Take set_initial_capacities_source as an example, the right-hand-side values of the source_constraints are set to the source_capacity.

for edge_metadata in values(graph.edge_data)
    (; edge) = edge_metadata
    if graph[edge...].subnetwork_id_source == subnetwork_id
        # If it is a source edge for this allocation problem
        if edge ∉ main_network_source_edges
            # Reset the source to the averaged flow over the last allocation period
            source_capacity = mean_input_flows[edge][]
            JuMP.set_normalized_rhs(
                source_constraints[edge],
                # It is assumed that the allocation procedure does not have to be differentiated.
                source_capacity,
            )
        end
    end
end

Apart from the set_initial_* function above, capacities of inlet are the allocated capacities from the main network to the subnetworks. Source constraints will be adapted based on the optimization type. This function is called separately and thus not part of the set_initial_values.

3.2 Looping over priorities

3.2.1 Updating capacities

While optimizing a given priority, the function set_capacities_flow_demand_outflow updates the constraints flow_demand_outflow. If the current priority is the same as the priority of the flow demand, constraints will be infinite, otherwise 0. At priorities where there is no flow demand, flow can go freely trough the node. When there is flow demand, flow is directed into the buffer. This is to make sure that flow can go to the node with the flow demand, even though the flow might have nowhere to go after that node.

The optimization objective function is updated based on the new demands and the given priority.

If a solution is found by the solver, the allocation result will be updated. And it will be saved, so the physical layer can make use of it.

Lastly, capacities and demands are updated, as shown below:

adjust_capacities_source!(allocation_model)
adjust_capacities_edge!(allocation_model)
adjust_capacities_basin!(allocation_model)
adjust_capacities_buffer!(allocation_model)
adjust_capacities_returnflow!(allocation_model, p)

for parameter in propertynames(p)
    demand_node = getfield(p, parameter)
    if demand_node isa AbstractDemandNode
        adjust_demands!(allocation_model, p, priority_idx, demand_node)
    end
end

3.3 Output data

The function save_demands_and_allocations saves the demand and the allocated value per demand node. And the function save_allocation_flows saves the optimized flows over the edges in the subnetwork. These values are saved in the record_demand and record_flow fields of the Allocation struct and only written to the output file at the end of the simulation.

3.4 Communicating to the physical layer

The function assign_allocations updates the subnetwork demand if the optimization task is collect_demands. It assigns the allocated amount to the UserDemand nodes with the result of the optimization if the optimization task is allocate. Afterwards, it writes the resulting flow to the Allocation object.

3.4.1 UserDemand abstraction

When allocation is active, the amount each UserDemand node is allowed to extract from its upstream basin is determined by the allocation algorithm. See here for more details on how allocation updates the UserDemand node.

3.4.2 Controlling pumps/weirs based on allocation results

N/A and TODO in this task.