Allocation

1 Introduction

This document provides a comprehensive technical reference for the allocation implementation in Ribasim. It bridges the gap between the mathematical formulation described in the concept documentation and the actual code implementation using JuMP.jl.

The allocation algorithm solves linear optimization problems to distribute water among competing demands. This document explains:

1. How mathematical formulations translate to JuMP code
2. Problem Building (placeholder initialization → real value updates)
3. Integration between the physical layer and optimization layer
4. Data structures and indexing patterns
5. Implementation details for each node type

2 Architecture Overview

2.1 Problem Building

The allocation optimization problems are build using a two-phase approach:

2.1.1 Phase 1: Problem Structure Initialization (allocation_init.jl)

During initialization, JuMP variables and constraints are created with placeholder values. These placeholders define the structure and relationships but not the actual values:

# Example from add_basin! function
current_storage = 1000.0 # Placeholder value
max_storage = 5000.0 # Placeholder value
problem[:basin_storage_change] = JuMP.@variable(
    problem,
    -current_storage / scaling.storage ≤
    basin_storage_change[node_id = basin_ids_subnetwork] ≤
    (max_storage - current_storage) / scaling.storage
)

Why placeholders? JuMP requires the problem structure (variables, constraints, objective) to be defined upfront. The actual physical values (water levels, flows, demands) are not yet known during initialization and will change at each allocation timestep.

2.1.2 Phase 2: Real Value Updates (allocation_optim.jl)

Before each optimization, the set_simulation_data! functions update constraints with real values from the physical layer:

# From set_simulation_data! for Basin
for basin_id in basin_ids_subnetwork
    basin_idx = basin_id.idx
    current_storage_basin = current_storage[basin_idx]

    # Update the variable bounds with actual storage values
    storage_change_variable = storage_change[basin_id]
    JuMP.set_lower_bound(
        storage_change_variable,
        -current_storage_basin / scaling.storage
    )
    # ... more updates
end

This pattern allows efficient repeated optimization: the problem structure remains fixed while only the constraint coefficients and bounds are updated.

2.2 Physical Layer Integration

The allocation layer operates on a linearized version of the physical layer. Here’s how physical quantities map to optimization variables:

2.2.1 Basin Levels and Storage

In the physical layer, basin level \(h\) is a nonlinear function of storage \(S\) determined by the basin profile (area-storage relationship). For optimization, we linearize around the current state:

Mathematical formulation: \[h^{n+1} \approx h^n + \frac{1}{A^n}(S^{n+1} - S^n)\]

where \(A^n\) is the basin area at the current timestep.

Code implementation:

# In set_simulation_data! for Basin
# Get linearization point
h_n = basin.current_level[basin_idx]
A_n = current_area[basin_idx]  # Area at current level
S_n = current_storage[basin_idx]

# Define the decision variable for storage change
# ΔS = S^{n+1} - S^n (scaled)
storage_change = problem[:basin_storage_change]

# The level at end of timestep is:
# h^{n+1} = h^n + (1/A^n) * ΔS
# This relationship is embedded in constraint coefficients (see below)

2.2.2 Connector Nodes (Flow as Function of Levels)

Nodes like TabulatedRatingCurve, LinearResistance, and ManningResistance have flow \(Q\) that depends on upstream level \(h_a\) and downstream level \(h_b\):

Mathematical formulation: \[Q^{n+1} \approx Q^n + \frac{\partial Q}{\partial h_a}(h_a - h_a^n) + \frac{\partial Q}{\partial h_b}(h_b - h_b^n)\]

For a basin downstream: \(h_b^{n+1} - h_b^n \approx \frac{1}{A_b}(S_b^{n+1} - S_b^n)\)

Code implementation:

# From linearize_connector_node!
# Evaluate flow and derivatives at current state
h_a = get_level(p, upstream_id, t)
h_b = get_level(p, downstream_id, t)
Q_n = flow_function(p, upstream_id, downstream_id, t_after)

# Compute partial derivatives
∂Q∂h_a = derivative_flow_wrt_level(p, upstream_id, downstream_id, t_after, :upstream)
∂Q∂h_b = derivative_flow_wrt_level(p, upstream_id, downstream_id, t_after, :downstream)

# Update constraint: Q = Q^n + ∂Q/∂h_a * (h_a - h_a^n) + ∂Q/∂h_b * (h_b - h_b^n)
# When h_b is a basin, substitute the linearized profile:
# Q = Q^n + (∂Q/∂h_b / A_b) * ΔS_b + ...

The function set_partial_derivative_wrt_level! handles converting level derivatives to storage derivatives:

function set_partial_derivative_wrt_level!(
    allocation_model::AllocationModel,
    node_id::NodeID,
    ∂q∂h::Float64,  # Partial derivative of flow w.r.t. level
    p::Parameters,
    constraint::JuMP.ConstraintRef,
)::Nothing
    (; problem, scaling) = allocation_model
    (; current_area) = p.state_and_time_dependent_cache

    # Convert ∂Q/∂h to ∂Q/∂S using ∂h/∂S = 1/A
    storage_change = problem[:basin_storage_change][node_id]
    JuMP.set_normalized_coefficient(
        constraint,
        storage_change,
        ∂q∂h / (current_area[node_id.idx] * scaling.flow * Δt_allocation)
    )
    return nothing
end

3 Mathematical Formulation to Code Mapping

3.1 Decision Variables

3.1.1 Flow Variables

Mathematical formulation: For each link \((i,j)\) in the network, a flow variable \(Q_{ij}\) represents the volumetric flow rate.

Code implementation:

# From add_flow!
problem[:flow] = JuMP.@variable(
    problem,
    flow_capacity_lower_bound(link, p_independent) / scaling.flow ≤
    flow[link = flow_links_subnetwork] ≤
    flow_capacity_upper_bound(link, p_independent) / scaling.flow
)

Key points: - Indexed by link::Tuple{NodeID, NodeID} representing (source, destination) - Bounds determined by node capacities (pump max flow, etc.) - Scaled by scaling.flow for numerical stability - Stored in sparse array for efficient access: problem[:flow][link]

3.1.2 Basin Storage Change Variables

Mathematical formulation: For each basin \(b\), a storage change variable \(\Delta S_b = S_b^{n+1} - S_b^n\) represents the change in storage over the allocation timestep.

Code implementation:

# From add_basin!
problem[:basin_storage_change] = JuMP.@variable(
    problem,
    -current_storage / scaling.storage ≤
    basin_storage_change[node_id = basin_ids_subnetwork] ≤
    (max_storage - current_storage) / scaling.storage
)

Key points: - Lower bound prevents storage from going negative - Upper bound prevents exceeding maximum basin capacity - Placeholder values replaced before each optimization - Scaled by scaling.storage for numerical conditioning

3.1.3 Allocated Flow Variables (Demand Nodes)

Mathematical formulation: For each demand node \(v\) and priority \(p \in P_v\): \[0 \le F^p_v \le d_v^p\]

Code implementation:

# From add_user_demand!
user_demand_allocated = JuMP.@variable(
    problem,
    0 ≤ user_demand_allocated[
        node_id = user_demand_ids_subnetwork,
        demand_priority = demand_priorities_all;
        has_demand_priority[node_id.idx, demand_priority_idx]
    ] ≤ d / scaling.flow
)

Key points: - Doubly-indexed by (node_id, demand_priority) - Only created for priorities where has_demand_priority is true - Bounded by demand value (placeholder, updated before optimization) - Similar structures for FlowDemand and LevelDemand

3.1.4 Error Variables

Mathematical formulation: For UserDemand/FlowDemand: \[E^p_v \ge 0, \quad \overline{E}^p_v \ge 0\]

Code implementation:

# From add_user_demand!
user_demand_error = JuMP.@variable(
    problem,
    user_demand_error[
        node_id = user_demand_ids_subnetwork,
        demand_priority = demand_priorities_all,
        objective_type = 1:2;  # 1 = absolute error, 2 = fairness error
        has_demand_priority[node_id.idx, demand_priority_idx]
    ] ≥ 0
)

Key points: - Triple-indexed: (node_id, demand_priority, objective_type) - objective_type = 1: absolute error \(E^p_v\) - objective_type = 2: fairness error \(\overline{E}^p_v\) - Non-negative to represent error magnitude

3.2 Constraints

3.2.1 Flow Conservation

Mathematical formulation: For conservative nodes (pumps, outlets, resistances), inflow equals outflow: \[Q_{\text{in}} = Q_{\text{out}}\]

Code implementation:

# From add_flow_conservation!
problem[Symbol(constraint_name)] = JuMP.@constraint(
    problem,
    [node_id = node_ids],
    flow[inflow_link[node_id.idx].link] == flow[outflow_link[node_id.idx].link],
    base_name = "flow_conservation_$node_name"
)

Key points: - Simple equality constraint between two flow variables - Applied to Pump, Outlet, LinearResistance, ManningResistance, TabulatedRatingCurve - Ensures mass conservation through the node

3.2.2 Volume Conservation (Basin Water Balance)

Mathematical formulation: \[\frac{dS}{dt} = \sum_{k=1}^{N_l} Q_k + f_{\text{pos}} - f_{\text{neg}}\]

Discretized with backward Euler: \[\frac{S^{n+1} - S^n}{\Delta t} = \sum_{k} Q_k^{n+1} + f^{n+1}_{\text{pos}} - f^{n+1}_{\text{neg}}\]

Or in terms of storage change \(\Delta S = S^{n+1} - S^n\): \[\Delta S = \Delta t \left(\sum_{k} Q_k^{n+1} + f^{n+1}_{\text{pos}} - f^{n+1}_{\text{neg}}\right)\]

Code implementation:

# From add_conservation!
problem[:volume_conservation] = JuMP.@constraint(
    problem,
    [node_id = basin_ids_subnetwork],
    storage_change[node_id] ==
        Δt_allocation / scaling.storage * (
            sum(flow[link] for link in inflow_links; init = 0.0) * scaling.flow -
            sum(
                low_storage_factor[node_id] * flow[link]
                for link in outflow_links;
                init = 0.0
            ) * scaling.flow +
            f_pos - f_neg  # Forcing terms (precipitation, evaporation, etc.)
        ),
    base_name = "volume_conservation"
)

Key points: - The low_storage_factor multiplies outflows to prevent negative storage - Forcing terms (f_pos, f_neg) represent precipitation, evaporation, etc. - Careful scaling to maintain numerical stability - Updated each optimization with current forcing values

3.2.3 UserDemand Allocated Sum

Mathematical formulation: The sum of allocated flows over all priorities equals the total inflow: \[F_{(b_{\text{upstream}}, v)} = \sum_{p \in P_v} F^p_v\]

Code implementation:

# From add_user_demand!
problem[:user_demand_allocated_sum_constraint] = JuMP.@constraint(
    problem,
    [node_id = user_demand_ids_subnetwork],
    flow[inflow_link[node_id.idx].link] ==
        sum(
            user_demand_allocated[node_id, demand_priority]
            for (demand_priority_idx, demand_priority) in
            enumerate(demand_priorities_all) if
            has_demand_priority[node_id.idx, demand_priority_idx];
            init = JuMP.AffExpr(0.0)
        ),
    base_name = "user_demand_allocated_sum"
)

Key points: - Ensures only demanded water enters UserDemand nodes - Sum is over priorities where the node has a demand - Links the flow variable to allocated flow variables

3.2.4 Error Constraints

Mathematical formulation: For UserDemand/FlowDemand: \[d_v^p \cdot E_v^p \ge d_v^p - F_v^p\]

Code implementation:

# From add_user_demand!
problem[:user_demand_relative_error_constraint] = JuMP.@constraint(
    problem,
    [
        node_id = user_demand_ids_subnetwork,
        demand_priority = demand_priorities_all,
        objective_type = 1:2;
        has_demand_priority[node_id.idx, demand_priority_idx]
    ],
    d * user_demand_error[node_id, demand_priority, objective_type] ≥
        d - user_demand_allocated[node_id, demand_priority],
    base_name = "user_demand_relative_error"
)

Key points: - Multiplication by demand d makes this effectively an absolute error - The error variable is forced to be at least the unmet demand fraction - Updated before each optimization with current demands

For LevelDemand:

Mathematical formulation: \[E^p_{b, \text{lower}} \ge s(h^p_{b, \min}) - (s(h_b^\text{init}) + \Delta S_b)\]

Code implementation:

# From add_level_demand!
problem[:storage_constraint_lower] = JuMP.@constraint(
    problem,
    [
        node_id = basin_ids_subnetwork_with_level_demand,
        demand_priority = demand_priorities_all;
        has_demand_priority[node_id.idx, demand_priority_idx]
    ],
    storage_change[node_id] +
        level_demand_error[node_id, demand_priority, 1] * current_area[node_id.idx] ≥
        (minimum_storage - starting_storage) / scaling.storage,
    base_name = "storage_constraint_lower"
)

Key points: - Error represents storage deficit below minimum level - Multiplication by area converts level error to storage error - Separate constraints for upper and lower bounds

3.2.5 Return Flow Constraints

Mathematical formulation: For UserDemand with return factor \(r_i(t)\): \[Q_{\text{out}} = r_i \cdot Q_{\text{in}}\]

Code implementation:

# From add_user_demand!
problem[:user_demand_return_flow] = JuMP.@constraint(
    problem,
    [node_id = user_demand_ids_subnetwork],
    flow[outflow_link[node_id.idx].link] ==
        return_factor * flow[inflow_link[node_id.idx].link],
    base_name = "user_demand_return_flow"
)

Key points: - return_factor is a placeholder, updated before optimization - Links inflow and outflow of UserDemand node - Models consumptive use (when return factor < 1)

3.2.6 Linearized Connector Node Constraints

Mathematical formulation: For nodes like TabulatedRatingCurve: \[Q^{n+1} \approx Q^n + \frac{\partial Q}{\partial h_a}(h_a^{n+1} - h_a^n) + \frac{\partial Q}{\partial h_b}(h_b^{n+1} - h_b^n)\]

Code implementation:

# From add_linearized_connector_node!
problem[Symbol(constraint_name)] = JuMP.@constraint(
    problem,
    [node_id = node_ids_subnetwork],
    flow[inflow_link[node_id.idx].link] ==
        q0 +  # Flow at linearization point
        (upstream_is_basin ?
            (∂q∂h_upstream / A_upstream) * storage_change[upstream_id] :
            0.0) +
        (downstream_is_basin ?
            (∂q∂h_downstream / A_downstream) * storage_change[downstream_id] :
            0.0),
    base_name = constraint_name
)

Key points: - q0, ∂q∂h_upstream, ∂q∂h_downstream are placeholders - Derivatives divided by area when basin is involved (converts \(\partial Q/\partial h\) to \(\partial Q/\partial S\)) - Updated in set_simulation_data! before each optimization - Handles various combinations (basin-basin, basin-boundary, etc.)

3.3 Objectives

The allocation optimization solves multiple objectives in sequence using lexicographic goal programming. Each objective is optimized, then the optimal value is constrained in subsequent objectives.

3.3.1 Primary Objective: Minimize Absolute Error Sum

Mathematical formulation: For flow demands (UserDemand, FlowDemand) at priority \(p\): \[\min \sum_{v,p} d_v^p \cdot E_v^p\]

For level demands at priority \(p\): \[\min \sum_{b; p\in P^\min_b} E^p_{b, \text{lower}} + \sum_{b; p\in P^\max_b} E^p_{b, \text{upper}}\]

Code implementation:

# From add_demand_objectives!
# For flow demands
objective_expression = JuMP.AffExpr(0.0)
for node_id in user_demand_ids_subnetwork
    if has_demand_priority[node_id.idx, demand_priority_idx]
        demand = get_demand(user_demand, node_id, demand_priority, t)
        JuMP.add_to_expression!(
            objective_expression,
            demand * user_demand_error[node_id, demand_priority, 1]
        )
    end
end

# Store for later optimization
push!(
    objective_metadata,
    ObjectiveMetadata(
        demand_priority,
        objective_expression,
        # ... other metadata
    )
)

Key points: - Error variables are weighted by demands (makes this absolute error minimization) - Separate objectives for each priority - Objectives are stored and optimized in order by optimize_multi_objective!

3.3.2 Secondary Objective: Minimize Fairness Error

Mathematical formulation: After minimizing the total error, minimize deviations from the average allocation rate.

For flow demands: \[\min \sum_{v,p} \overline{E}_v^p\]

where \(\overline{E}_v^p \ge E_v^p - G^p\) and \(G^p\) is the global average allocation rate.

Code implementation:

# From add_demand_objectives!
# Compute average error (constraint on average error variable)
average_flow_unit_error_constraint =
    @constraint(
        problem,
        [demand_priority = demand_priorities_all; ...],
        sum_demands * average_flow_unit_error[demand_priority] ==
            sum(d * error[node_id, demand_priority, 1] for node_id, d in demands),
        base_name = "average_flow_unit_error"
    )

# Fairness error: overline_E >= E - G
# (Implicitly defined through constraints)

# Objective sums fairness errors
objective_expression_fairness = sum(
    user_demand_error[node_id, demand_priority, 2]
    for node_id in user_demand_ids_subnetwork if ...
)

Key points: - average_flow_unit_error represents \(G^p\) - Fairness errors (objective_type = 2) penalize being worse than average - Optimized after the primary objective is satisfied

3.3.3 Route Priority Objective

Mathematical formulation: \[\min \sum w_i \cdot Q_i\]

where \(w_i\) is the route priority weight (cost) for node \(i\).

Code implementation:

# From add_route_priority_objective!
objective_expression = JuMP.AffExpr(0.0)
for link in flow_links
    # Get the route priority weight for the source node
    weight = route_priority_weight[link[1]]
    if weight > 0
        JuMP.add_to_expression!(
            objective_expression,
            weight * flow[link]
        )
    end
end

Key points: - Optimized after all demand objectives are satisfied - Only affects routing, not how much is allocated - Higher weight = less preferred route (higher cost)

3.3.4 Low Storage Factor Objective

Mathematical formulation: \[\max \sum_{b \in \text{basins}} \alpha_b\]

where \(\alpha_b\) is the low storage factor for basin \(b\).

Code implementation:

# From add_low_storage_factor_objective!
objective_expression = sum(
    low_storage_factor[node_id]
    for node_id in basin_ids_subnetwork
)

# Note: Maximizing low_storage_factor is equivalent to minimizing its negative
push!(
    objective_metadata,
    ObjectiveMetadata(
        Int32(-1),  # Special priority for this objective
        -objective_expression,  # Negative for maximization
        # ...
    )
)

Key points: - Maximizes the low storage factor to allow more outflow from basins - Optimized last to avoid infeasibility from emptying basins - Special priority (-1) ensures it runs after demand objectives

4 Updating Constraints Before Optimization

Before each optimization run, constraints must be updated with current values from the physical layer. This is done through a series of set_simulation_data! functions.

4.1 Basin Updates

function set_simulation_data!(
    allocation_model::AllocationModel,
    basin::Basin,
    p::Parameters,
    t::Float64,
)::Bool
    (; problem, node_ids_in_subnetwork, scaling) = allocation_model
    (; basin_ids_subnetwork) = node_ids_in_subnetwork
    storage_change = problem[:basin_storage_change]

    for basin_id in basin_ids_subnetwork
        basin_idx = basin_id.idx

        # Get current state from physical layer
        current_storage_basin = current_storage[basin_idx]
        max_storage_basin = basin.storage_to_level[basin_idx].t[end]

        # Update bounds on storage change variable
        JuMP.set_lower_bound(
            storage_change[basin_id],
            -current_storage_basin / scaling.storage
        )
        JuMP.set_upper_bound(
            storage_change[basin_id],
            (max_storage_basin - current_storage_basin) / scaling.storage
        )

        # Update volume conservation constraint with current forcing
        # (precipitation, evaporation, etc.)
        # ...
    end
    return errors
end

Key concepts: - Variable bounds are updated each timestep based on current basin storage - Forcing terms (precipitation, evaporation) are computed and added to constraints - Errors are detected (e.g., storage outside valid range) and reported

4.2 Connector Node Updates (Linearization)

function linearize_connector_node!(
    allocation_model::AllocationModel,
    connector_node::AbstractParameterNode,
    flow_constraint,
    flow_function::Function,
    p::Parameters,
    t::Float64,
)
    (; scaling, Δt_allocation) = allocation_model

    # Evaluate at the end of the allocation timestep (backward Euler)
    t_after = t + Δt_allocation

    for node_id in only(flow_constraint.axes)
        # Get upstream and downstream IDs
        upstream_id = connector_node.inflow_link[node_id.idx].link[1]
        downstream_id = connector_node.outflow_link[node_id.idx].link[2]

        # Get levels at linearization point
        h_a = get_level(p, upstream_id, t)
        h_b = get_level(p, downstream_id, t)

        # Evaluate flow at linearization point
        Q_n = flow_function(p, upstream_id, downstream_id, t_after)

        # Compute partial derivatives
        ∂Q∂h_a = compute_derivative(p, upstream_id, downstream_id, t_after, :upstream)
        ∂Q∂h_b = compute_derivative(p, upstream_id, downstream_id, t_after, :downstream)

        # Update constraint with linearization
        constraint_ref = flow_constraint[node_id]

        # Set constant term (Q_n - ∂Q/∂h_a * h_a - ∂Q/∂h_b * h_b)
        JuMP.set_normalized_rhs(
            constraint_ref,
            (Q_n - ∂Q∂h_a * h_a - ∂Q∂h_b * h_b) / scaling.flow
        )

        # Set coefficients for basin storage variables if applicable
        if upstream_id.type == NodeType.Basin
            set_partial_derivative_wrt_level!(
                allocation_model, upstream_id, ∂Q∂h_a, p, constraint_ref
            )
        end
        if downstream_id.type == NodeType.Basin
            set_partial_derivative_wrt_level!(
                allocation_model, downstream_id, ∂Q∂h_b, p, constraint_ref
            )
        end
    end
end

Key concepts: - Linearization is performed at the current physical state - Derivatives are computed using physical layer functions - Constraint coefficients are updated to match the Taylor expansion - Backward Euler: evaluate at end of timestep (\(t + \Delta t\))

4.3 Demand Updates

function set_demands!(
    allocation_model::AllocationModel,
    node::Union{UserDemand, FlowDemand},
    # ... other parameters
)::Nothing
    for (demand_priority_idx, demand_priority) in enumerate(demand_priorities_all)
        for node_id in demand_node_ids_subnetwork
            if !has_demand_priority[node_id.idx, demand_priority_idx]
                continue
            end

            # Get current demand from physical layer
            demand = get_demand(node, node_id, demand_priority, t)

            # Update error constraint with current demand
            error_constraint = node_relative_error_constraint[node_id, demand_priority, 1]
            JuMP.set_normalized_coefficient(
                error_constraint,
                node_error[node_id, demand_priority, 1],
                demand / scaling.flow
            )

            # Update upper bound on allocated variable
            JuMP.set_upper_bound(
                node_allocated[node_id, demand_priority],
                demand / scaling.flow
            )

            # Similar updates for fairness error constraints...
        end
    end
    return nothing
end

Key concepts: - Demands are time-varying and interpolated from input tables - Constraint coefficients are updated to reflect current demands - Upper bounds on allocation variables ensure allocated ≤ demanded

5 The Optimization Loop

The main optimization loop in optimize_multi_objective! solves objectives in sequence:

function optimize_multi_objective!(
    allocation_model::AllocationModel,
    # ...
)::Nothing
    (; problem, objectives, temporary_constraints) = allocation_model

    for metadata in objectives.objective_metadata
        (; demand_priority, objective_expression) = metadata

        # Set the current objective
        JuMP.@objective(problem, Min, objective_expression)

        # Solve
        JuMP.optimize!(problem)

        # Check termination status
        if JuMP.termination_status(problem) != JuMP.OPTIMAL
            # Handle infeasibility...
        end

        # Add constraint to preserve this objective's value
        optimal_value = JuMP.objective_value(problem)
        constraint = JuMP.@constraint(
            problem,
            objective_expression <= optimal_value + tolerance
        )
        push!(temporary_constraints, constraint)
    end

    # Clean up temporary constraints after optimization
    delete_temporary_constraints!(allocation_model)

    return nothing
end

Key concepts: - Lexicographic optimization: Each objective is optimized in order - After optimizing objective \(i\), add constraint: \(\text{obj}_i \le \text{optimal}_i + \epsilon\) - This ensures later objectives don’t degrade earlier ones - Temporary constraints are removed after all objectives are solved - If any objective is infeasible, the entire optimization fails

6 Scaling for Numerical Stability

The optimization uses scaling factors to improve numerical conditioning:

struct ScalingFactors
    storage::Float64  # Typical storage value (m³)
    flow::Float64     # Typical flow rate (m³/s)
end

function ScalingFactors(
    p_independent::ParametersIndependent,
    subnetwork_id::Int32,
    Δt_allocation::Float64,
)
    # Compute average half-storage across basins
    max_storages = [
        basin.storage_to_level[node_id.idx].t[end]
        for node_id in basin.node_id
        if graph[node_id].subnetwork_id == subnetwork_id
    ]
    mean_half_storage = sum(max_storages) / (2 * length(max_storages))

    return ScalingFactors(
        storage = mean_half_storage,
        flow = mean_half_storage / Δt_allocation
    )
end

Key concepts: - Storage variables are divided by scaling.storage - Flow variables are divided by scaling.flow - Typical values are O(1) in the scaled problem - Improves numerical stability and solver performance - Must scale/unscale when communicating with physical layer

6.1 Example: Scaled Volume Conservation

Unscaled formulation: \[\Delta S \text{ [m³]} = \Delta t \text{ [s]} \cdot \left( Q_{\text{in}} - Q_{\text{out}} \text{ [m³/s]} \right)\]

Scaled formulation: \[\tilde{\Delta S} \cdot S_{\text{scale}} = \Delta t \cdot \left( \tilde{Q}_{\text{in}} \cdot Q_{\text{scale}} - \tilde{Q}_{\text{out}} \cdot Q_{\text{scale}} \right)\]

where \(\tilde{\Delta S} = \Delta S / S_{\text{scale}}\) and \(\tilde{Q} = Q / Q_{\text{scale}}\).

Dividing through by \(S_{\text{scale}}\): \[\tilde{\Delta S} = \frac{\Delta t \cdot Q_{\text{scale}}}{S_{\text{scale}}} \left( \tilde{Q}_{\text{in}} - \tilde{Q}_{\text{out}} \right)\]

In the code:

storage_change[node_id] ==
    Δt_allocation / scaling.storage * (
        sum(flow[link] for link in inflow_links) * scaling.flow - ...
    )

7 Data Structures and Indexing

7.1 NodeID

Nodes are identified by the NodeID struct:

struct NodeID
    type::NodeType.T  # e.g., NodeType.Basin, NodeType.UserDemand
    idx::Int32        # Index into the type-specific data arrays
    value::Int32      # User-facing ID from input file
end

Key concepts: - idx is used to index into type-specific arrays (e.g., basin.storage[basin_id.idx]) - value is the ID from the input file (for error messages, output) - type distinguishes different node types

7.2 Links

Links are represented as tuples of NodeIDs:

link::Tuple{NodeID, NodeID}  # (source, destination)

Key concepts: - Used to index flow variables: problem[:flow][link] - Direction matters: (A, B) ≠ (B, A) - Stored in graph edge metadata

7.3 JuMP Sparse Arrays

JuMP uses SparseAxisArray for multi-dimensional variables:

# Example: doubly-indexed variable
user_demand_allocated[node_id, demand_priority]

# Example: triply-indexed variable
user_demand_error[node_id, demand_priority, objective_type]

Key concepts: - Only creates variables for valid index combinations - Filtered by conditions (e.g., has_demand_priority[node_id.idx, demand_priority_idx]) - Efficient for sparse problems (not all nodes have all priorities)

7.4 AllocationModel Struct

The main data structure:

struct AllocationModel
    subnetwork_id::Int32
    problem::JuMP.Model  # The JuMP optimization problem
    node_ids_in_subnetwork::NodeIDsInSubnetwork
    scaling::ScalingFactors
    Δt_allocation::Float64

    # Cumulative tracking for output
    cumulative_realized_volume::Dict{Tuple{NodeID, NodeID}, Float64}
    cumulative_boundary_volume::Dict{Tuple{NodeID, NodeID}, Float64}

    # Forcing volumes (precipitation, evaporation)
    explicit_positive_forcing_volume::Dict{NodeID, Float64}
    implicit_negative_forcing_volume::Dict{NodeID, Float64}

    # Optimization objectives
    objectives::AllocationObjectives

    # Temporary constraints for lexicographic optimization
    temporary_constraints::Vector{JuMP.ConstraintRef}

    # For routing optimization
    route_priority_expression::JuMP.AffExpr
end

8 Secondary Networks and Primary-Secondary Connections

8.1 Primary Network

The primary network (subnetwork_id = 1) represents the main water system. It:

• Contains its own demand nodes
• Connects to secondary networks via Pump or Outlet nodes
• Treats each secondary network as a single demand node

8.2 Secondary Networks

Secondary networks (subnetwork_id > 1):

• Can only connect to the primary network
• Cannot connect to each other directly
• Have their own internal allocation optimization

8.3 The Two-Stage Process

8.3.1 Stage 1: Demand Collection

For each secondary network:

1. Set inflow from primary network to unlimited (temporarily)
2. Solve allocation to determine total unmet demand
3. This demand becomes the secondary network’s “request” to the primary network

function preprocess_demand_collection!(
    allocation_model::AllocationModel,
    p_independent::ParametersIndependent,
)::Nothing
    # Allow unlimited inflow from primary network
    for link in primary_network_connections[subnetwork_id]
        flow_var = flow[link]
        JuMP.set_upper_bound(flow_var, MAX_ABS_FLOW / scaling.flow)
        JuMP.set_lower_bound(flow_var, 0)
    end

    # Solve to find total demand...
end

8.3.2 Stage 2: Allocation

In primary network:

1. Solve allocation considering:
   • Primary network’s own demands
   • Aggregated demands from secondary networks
2. Allocated amounts to secondary networks determine their inflow limits

In each secondary network:

1. Set inflow constraint to the allocated amount from primary network
2. Solve allocation to distribute water to internal demands

function allocate_flows_to_subnetwork(
    allocation_models::Vector{AllocationModel},
    primary_network_connections,
)::Nothing
    primary_network = get_primary_network(allocation_models)
    primary_problem = primary_network.problem

    for secondary_network in get_secondary_networks(allocation_models)
        # Get allocated flow from primary network
        for link in primary_network_connections[secondary_network.subnetwork_id]
            allocated_flow = JuMP.value(primary_problem[:flow][link])

            # Set as capacity for secondary network
            secondary_problem = secondary_network.problem
            JuMP.set_upper_bound(
                secondary_problem[:flow][link],
                allocated_flow
            )
        end
    end
end

9 Warm Start

To improve solver performance, the optimization can be warm-started with flow rates from the physical layer:

function warm_start!(allocation_model::AllocationModel, integrator::DEIntegrator)::Nothing
    (; problem, scaling, Δt_allocation) = allocation_model
    flow = problem[:flow]
    du = get_du(integrator)  # du/dt from physical layer

    # Extrapolate current flow rates
    for link in only(flow.axes)
        # Get current flow from physical layer
        current_flow = get_flow(du, p, link)

        # Set as starting value for optimization
        JuMP.set_start_value(
            flow[link],
            current_flow / scaling.flow
        )
    end

    # Similar for basin storage changes...
end

Key concepts: - Starting values guide the solver to a good initial solution - Can significantly reduce solver iterations - Based on extrapolating current physical state forward

10 Output and Communication with Physical Layer

10.1 Parsing Results

After optimization, results must be extracted and stored:

function parse_allocations!(
    integrator::DEIntegrator,
    allocation_model::AllocationModel,
)::Nothing
    (; problem, subnetwork_id, scaling) = allocation_model

    # Extract allocated amounts per demand node per priority
    for node_id in user_demand_ids_subnetwork
        for (demand_priority_idx, demand_priority) in enumerate(demand_priorities_all)
            if !has_demand_priority[node_id.idx, demand_priority_idx]
                continue
            end

            # Get allocated amount from optimization result
            allocated = JuMP.value(user_demand_allocated[node_id, demand_priority])
            allocated_unscaled = allocated * scaling.flow

            # Store for output
            record_demand[subnetwork_id][demand_priority][node_id] = (
                demand = current_demand,
                allocated = allocated_unscaled,
                realized = 0.0  # To be filled in after physical layer runs
            )
        end
    end
end

10.2 Applying Allocation Results

Allocated amounts must be communicated to the physical layer:

10.2.1 UserDemand Nodes

# In the physical layer, UserDemand nodes have a max_flow_rate that limits extraction
function update_user_demand!(user_demand::UserDemand, allocation_result)
    for node_id in user_demand.node_id
        # Get allocated amount summed over all priorities
        total_allocated = sum(
            allocation_result[node_id][priority].allocated
            for priority in priorities
        )

        # Set as maximum extraction rate
        user_demand.max_flow_rate[node_id.idx] = total_allocated
    end
end

10.2.2 Pump/Outlet Control

When Pump or Outlet has control state Ribasim.allocation:

function apply_control_from_allocation!(
    node::Union{Pump, Outlet},
    allocation_model::AllocationModel,
    integrator::DEIntegrator,
)::Nothing
    (; problem, scaling) = allocation_model
    flow = problem[:flow]

    for node_id in controlled_node_ids
        # Get optimized flow from allocation
        link = (inflow_id(node_id), outflow_id(node_id))
        optimized_flow = JuMP.value(flow[link]) * scaling.flow

        # Set as target flow rate in physical layer
        node.flow_rate[node_id.idx] = optimized_flow
    end
end

11 Handling Infeasibility

When the optimization problem is infeasible, diagnostic tools help identify the cause:

11.1 Infeasibility Analysis

function analyze_infeasibility(
    allocation_model::AllocationModel,
    t::Float64,
    config::Config,
)::JuMP.TerminationStatusCode
    (; problem) = allocation_model

    # Find Irreducible Inconsistent Subsystem (IIS)
    data_infeasibility = MathOptAnalyzer.analyze(
        MathOptAnalyzer.Infeasibility.Analyzer(),
        problem;
        optimizer = get_optimizer(),
    )

    # Extract conflicting constraints
    violated_constraints = [...]

    # Try relaxing constraints to identify issues
    constraint_to_penalty = Dict(
        constraint => (isempty(JuMP.name(constraint)) ? 1.0 : 0.5)
        for constraint in violated_constraints
    )
    constraint_to_slack = JuMP.relax_with_penalty!(problem, constraint_to_penalty)
    JuMP.optimize!(problem)

    # Report which constraints are in conflict
    for irreducible_infeasible_subset in data_infeasibility.iis
        constraint_violations = [...]
        @error "Set of incompatible constraints found" constraint_violations
    end

    return JuMP.INFEASIBLE
end

Key concepts: - IIS = minimal set of constraints that cannot be satisfied simultaneously - Constraint relaxation helps identify which constraints are problematic - Named constraints are more informative for debugging

11.2 Common Infeasibility Causes

1. Insufficient water supply: Total demand exceeds total available water
2. Conflicting level demands: Min/max levels cannot be satisfied simultaneously
3. Capacity constraints: Network capacity insufficient to meet demands
4. Storage constraints: Basin storage limits prevent meeting level demands

12 Performance Considerations

12.1 Problem Size

• Variables: O(links + basins + demand_nodes × priorities)
• Constraints: O(links + basins + demand_nodes × priorities)
• Objectives: O(priorities)

For a network with 100 nodes, 200 links, 20 demand nodes, and 4 priorities: - ~300 variables - ~400 constraints - 4-8 objectives (depending on objectives configured)

12.2 Solver Performance

The HiGHS solver is configured for allocation problems:

function get_optimizer()
    JuMP.optimizer_with_attributes(
        HiGHS.Optimizer,
        "log_to_console" => false,
        "time_limit" => 60.0,
        "random_seed" => 0,
        "small_matrix_value" => 1e-12,  # Numerical threshold
    )
end

Tips for performance: - Use scaling to keep values O(1) - Minimize number of priorities (each adds objectives) - Warm start with physical layer flows - Simplify network topology where possible

12.3 Allocation Timestep Selection

The allocation timestep \(\Delta t_{\text{allocation}}\) must be chosen carefully:

• Too small: Frequent optimization overhead, linearization always valid
• Too large: Linearization may be inaccurate, less responsive to changes
• Typical range: 1 day to 1 week for water resources models

13 Debugging Tips

13.1 Inspecting the Problem

# Write problem to file for manual inspection
write_problem_to_file(problem, config)

# Check variable values
println(JuMP.value(flow[link]))
println(JuMP.value(storage_change[basin_id]))

# Check constraint values
println(JuMP.normalized_rhs(constraint))
println(JuMP.normalized_coefficient(constraint, variable))

13.2 Common Issues

13.2.1 Variables at Bounds

function get_bounds_hit(variable::JuMP.VariableRef)::Tuple{Bool, Bool}
    hit_lower_bound = if JuMP.has_lower_bound(variable)
        JuMP.value(variable) ≤ JuMP.lower_bound(variable)
    else
        false
    end

    hit_upper_bound = if JuMP.has_upper_bound(variable)
        JuMP.value(variable) ≥ JuMP.upper_bound(variable)
    else
        false
    end

    return hit_lower_bound, hit_upper_bound
end

If many variables are at bounds, the problem may be under-constrained or infeasible.

13.2.2 Numerical Scaling Issues

function analyze_scaling(
    allocation_model::AllocationModel,
    t::Float64,
    config::Config,
)::Nothing
    data_numerical = MathOptAnalyzer.analyze(
        MathOptAnalyzer.Numerical.Analyzer(),
        problem;
        threshold_small = 1e-12,
        threshold_large = 1e6,
    )

    # Check for poorly scaled coefficients
    for data in data_numerical.matrix_small
        @error "Too small coefficient" data.coefficient
    end

    for data in data_numerical.matrix_large
        @error "Too large coefficient" data.coefficient
    end
end

14 Summary

This technical reference provides the link between mathematical formulation and code implementation:

1. Two-phase approach: Structure defined at initialization with placeholders, values updated before each optimization
2. Physical layer integration: Linearization of basin profiles and connector nodes around current state
3. Decision variables: Flows, storage changes, allocated amounts, errors
4. Constraints: Flow conservation, volume conservation, demand relationships, linearized physics
5. Objectives: Lexicographic goal programming for demands, fairness, route priorities
6. Scaling: Improves numerical stability by keeping values O(1)
7. Data structures: NodeID-based indexing, sparse JuMP arrays, link tuples
8. Primary-secondary networks: Two-stage allocation process

For user-facing documentation, see Allocation Concept. For implementation code, see core/src/allocation_init.jl, core/src/allocation_optim.jl, and core/src/allocation_util.jl.