Review comments adressed

core/src/Ribasim.jl

@@ -17,6 +17,7 @@ module Ribasim
 import IterTools
 import BasicModelInterface as BMI
 import HiGHS
+import JuMP.Model as JuMPModel
 import TranscodingStreams
 
 using Arrow: Arrow, Table
@@ -32,17 +33,17 @@ using ForwardDiff: pickchunksize
 using DiffEqCallbacks
 using Graphs:
     add_edge!,
-    rem_edge!,
     adjacency_matrix,
     all_neighbors,
     DiGraph,
+    Edge,
     edges,
     inneighbors,
     nv,
     outneighbors,
-    SimpleEdge
+    rem_edge!
 
-using JuMP
+using JuMP: @variable, @constraint, @objective, set_normalized_rhs, optimize!, value
 using Legolas: Legolas, @schema, @version, validate, SchemaVersion, declared
 using Logging: current_logger, min_enabled_level, with_logger
 using LoggingExtras: EarlyFilteredLogger, LevelOverrideLogger

core/src/create.jl

@@ -231,7 +231,7 @@ function Connectivity(db::DB, config::Config, chunk_size::Int)::Connectivity
     end
 
     # TODO: Create allocation models from input here
-    allocation_model = AllocationModel[]
+    allocation_models = AllocationModel[]
 
     return Connectivity(
         graph_flow,
@@ -242,7 +242,7 @@ function Connectivity(db::DB, config::Config, chunk_size::Int)::Connectivity
         edge_ids_control,
         edge_connection_types_flow,
         edge_connection_types_control,
-        allocation_model,
+        allocation_models,
     )
 end
 

core/src/solve.jl

@@ -6,14 +6,13 @@ const VectorInterpolation =
 
 """
 Store information for a subnetwork used for allocation.
-For more information see allocation.jl.
 
 node_id: All the IDs of the nodes that are in this subnetwork
 node_id_mapping: Mapping Dictionary; model_node_id => AG_node_id where such a correspondence exists
     (all AG node ids are in the values)
-node_id_mapping_inverse: The inverse of node_id_mapping, Dictionary; AG_node_id => model_node_id
-Source edge mapping: AG source node ID => subnetwork source edge id
-graph_max_flow: The graph used for the allocation problems
+node_id_mapping_inverse: The inverse of node_id_mapping, Dictionary; AG node ID => model node ID
+Source edge mapping: AG source node ID => subnetwork source edge ID
+graph_allocation: The graph used for the allocation problems
 capacity: The capacity per edge of the allocation graph, as constrained by nodes that have a max_flow_rate
 model: The JuMP.jl model for solving the allocation problem
 Δt_allocation: The time interval between consecutive allocation solves
@@ -25,7 +24,7 @@ struct AllocationModel
     source_edge_mapping::Dict{Int, Int}
     graph_allocation::DiGraph{Int}
     capacity::SparseMatrixCSC{Float64, Int}
-    model::JuMP.Model
+    model::JuMPModel
     Δt_allocation::Float64
 end
 
@@ -52,7 +51,7 @@ struct Connectivity{T}
     edge_ids_control::Dictionary{Tuple{Int, Int}, Int}
     edge_connection_type_flow::Dictionary{Int, Tuple{Symbol, Symbol}}
     edge_connection_type_control::Dictionary{Int, Tuple{Symbol, Symbol}}
-    allocation_model::Vector{AllocationModel}
+    allocation_models::Vector{AllocationModel}
     function Connectivity(
         graph_flow,
         graph_control,

core/test/allocation.jl

@@ -21,7 +21,7 @@ using JuMP: value
     t = 0.0
 
     allocation_model =
-        Ribasim.get_allocation_model(p, subgraph_node_ids, source_edge_ids, Δt_allocation)
+        Ribasim.AllocationModel(p, subgraph_node_ids, source_edge_ids, Δt_allocation)
     Ribasim.allocate!(p, allocation_model, t)
 
     F = value.(allocation_model.model[:F])

docs/core/allocation.qmd

@@ -2,9 +2,9 @@
 title: "Allocation"
 ---
 
-Allocation is the process of associating an allocated abstraction flow rate to user nodes in the model based on information about sources, user demands over various priorities, constraints introduced by nodes, local water availability and graph topology. The allocation procedure implemented in Ribasim is heavily inspired by the [maximum flow problem](https://en.wikipedia.org/wiki/Maximum_flow_problem).
+Allocation is the process of assigning an allocated abstraction flow rate to user nodes in the model based on information about sources, user demands over various priorities, constraints introduced by nodes, local water availability and graph topology. The allocation procedure implemented in Ribasim is heavily inspired by the [maximum flow problem](https://en.wikipedia.org/wiki/Maximum_flow_problem).
 
-The allocation problem is solved per subnetwork of the Ribasim model. The subnetwork is used to formulate an optimization problem which is solved using the `JuMP.jl` package.
+The allocation problem is solved per subnetwork of the Ribasim model. The subnetwork is used to formulate an optimization problem with the [JuMP](https://jump.dev/JuMP.jl/stable/) package, which is solved using the [HiGHS solver](https://highs.dev/). See also the example of solving the maximum flow problem with `JuMP.jl` [here](https://jump.dev/JuMP.jl/stable/tutorials/linear/network_flows/#The-max-flow-problem).
 
 # The allocation problem
 
@@ -18,7 +18,7 @@ The allocation problem is solved per subgraph, where a subgraph is given by a su
 
 ### Source flows
 
-Sources are indicated a set of edges in the subnetwork
+Sources are indicated by a set of edges in the subnetwork
 $$
 E_S^\text{source} \subset \left(S \times S\right) \cap E.
 $$
@@ -37,7 +37,7 @@ On this page we assume that the priorities are given by all integers from $1$ to
 
 ### Vertical fluxes and local storage
 
-Apart from the source flows denoted by edges, there are other sources of water in the subnetwork, associated with the basins in the subnetwork $B_S = B \cap S$. Firstly there is the sum of the horizontal fluxes (precipitation, evaporation, infiltration and drainage) for each basin
+Apart from the source flows denoted by edges, there are other sources of water in the subnetwork, associated with the basins in the subnetwork $B_S = B \cap S$. Firstly there is the sum of the vertical fluxes (precipitation, evaporation, infiltration and drainage) for each basin
 $$
     \phi_i(t), \quad \forall i \in B_S.
 $$
@@ -62,7 +62,7 @@ Both fractional flow nodes and user nodes dictate proportional relationships bet
 
 A new graph is created from the subnetwork, which we call the allocation graph. To indicate the difference between subnetwork data and allocation graph data, the allocation graph data is denoted with a hat. The allocation graph consists of:
 
-- Nodes $\hat{V_S}$, where each basin, source and user in the subnetwork get a node in the allocation graph. Also nodes that have fractional flow outneighbors get a node in the allocation graph. The implementation makes heavy use of the node id mapping $m_S : i \mapsto \hat{i}$ to translate from subnetwork node Ds to allocation graph node IDs. Unless specified otherwise, we assume this relationship between index symbols that appear both with and without a hat.
+- Nodes $\hat{V_S}$, where each basin, source and user in the subnetwork get a node in the allocation graph. Also nodes that have fractional flow outneighbors get a node in the allocation graph. The implementation makes heavy use of the node id mapping $m_S : i \mapsto \hat{i}$ to translate from subnetwork node IDs to allocation graph node IDs. Unless specified otherwise, we assume this relationship between index symbols that appear both with and without a hat.
 - Edges $\hat{E_S}$, where the edges in the allocation graph are given by one or more edges in the subnetwork, where those edges connect nodes in the subnetwork that have an equivalent in the allocation graph. The direction of the edges in the allocation graph is given by the direction constraints in the subnetwork.
 
 For notational convenience, we use the notation
@@ -76,9 +76,9 @@ for the set of in-neighbors and out-neighbors of a node in the allocation graph
 
 ### The allocation graph capacities
 
-The capacities of the edges of the allocation graph are collected in the sparse capacity matrix $\hat{C}_S \in \overline{\mathbb{R}}_{\ge 0}^{\hat{n}\times\hat{n}}$ where $\hat{n}$ is the number of nodes in the allocation graph. The capacities can be infinity.
+The capacities of the edges of the allocation graph are collected in the sparse capacity matrix $\hat{C}_S \in \overline{\mathbb{R}}_{\ge 0}^{\hat{n}\times\hat{n}}$ where $\hat{n}$ is the number of nodes in the allocation graph. The capacities can be infinite.
 
-The The capacities are determined in 3 different ways:
+The capacities are determined in 3 different ways:
 
 - If an edge does not exist, i.e. $(\hat{i},\hat{j}) \notin \hat{E}$ for certain $1 \le \hat{i},\hat{j}\le \hat{n}$, then $(\hat{C}_S)_{\hat{i},\hat{j}} = 0$;
 - The capacity of the edge $\hat{e} \in \hat{E_S}$ is given by the smallest `max_flow_rate` of the nodes along the equivalent edges in the subnetwork. If there are no nodes with a `max_flow_rate`, the edge capacity is infinite;
@@ -165,14 +165,59 @@ $$ {#eq-fractionalflowconstraint}
 
 - Flow sign: Furthermore there are the non-negativity constraints for the flows and allocations, see [The optimization variables](allocation.qmd#the-optimization-variables).
 
-:::{.callout-note}
-Note that $\Phi_{\hat{k}}$ can become negative, in which case @eq-flowconservationconstraint becomes impossible to satisfy. What to do in this case?
-:::
-
 ## Final notes on the allocation problem
 
 ### Users using their own return flow
 
-If not explicitly avoided, users can use their own return flow in this allocation problem formulation. Therefore, Return flow of users is only taken into account by allocation if that return flow is downstream of the user where it comes from. That is, if there is no path in the directed allocation graph from the user outflow node back to the user.
+If not explicitly avoided, users can use their own return flow in this allocation problem formulation.
+Therefore, return flow of users is only taken into account by allocation if that return flow is downstream of the user where it comes from. That is, if there is no path in the directed allocation graph from the user outflow node back to the user.
 
 # Solving the allocation problem
+
+The text below shows a representation generated by `JuMP.jl` of an optimization as described in the previous section.
+
+```
+Max A_basin[2] + A_basin[3] + A_basin[5] + 0.5 A_user_4[1] + 0.25 A_user_4[2] + 0.125 A_user_4[3] + 0.5 A_user_6[1] + 0.25 A_user_6[2] + 0.125 A_user_6[3] + 0.5 A_user_1[1] + 0.25 A_user_1[2] + 0.125 A_user_1[3]
+Subject to
+ allocation_sum[1] : -F[1] + A_user_1[1] + A_user_1[2] + A_user_1[3] == 0
+ allocation_sum[4] : -F[3] + A_user_4[1] + A_user_4[2] + A_user_4[3] == 0
+ allocation_sum[6] : -F[5] + A_user_6[1] + A_user_6[2] + A_user_6[3] == 0
+ A_user_1[1] >= 0
+ A_user_1[2] >= 0
+ A_user_1[3] >= 0
+ A_user_4[1] >= 0
+ A_user_4[2] >= 0
+ A_user_4[3] >= 0
+ A_user_6[1] >= 0
+ A_user_6[2] >= 0
+ A_user_6[3] >= 0
+ demand_user_1[1] : A_user_1[1] <= 4
+ demand_user_1[2] : A_user_1[2] <= 0
+ demand_user_1[3] : A_user_1[3] <= 0
+ demand_user_4[1] : A_user_4[1] <= 0
+ demand_user_4[2] : A_user_4[2] <= 2
+ demand_user_4[3] : A_user_4[3] <= 0
+ demand_user_6[1] : A_user_6[1] <= 0
+ demand_user_6[2] : A_user_6[2] <= 1
+ demand_user_6[3] : A_user_6[3] <= 0
+ basin_allocation[2] : A_basin[2] <= 0
+ basin_allocation[3] : A_basin[3] <= 0
+ basin_allocation[5] : A_basin[5] <= 0
+ capacity[2] : F[2] <= 3
+ capacity[4] : F[4] <= 4
+ source[7] : F[6] <= 4.5
+ flow_conservation[2] : F[1] - F[2] <= 0
+ flow_conservation[3] : F[2] + F[3] - F[4] <= 0
+ flow_conservation[5] : F[4] + F[5] - F[6] <= 0
+ F[1] >= 0
+ F[2] >= 0
+ F[3] >= 0
+ F[4] >= 0
+ F[5] >= 0
+ F[6] >= 0
+ A_basin[2] >= 0
+ A_basin[3] >= 0
+ A_basin[5] >= 0
+```
+
+A more detailed explanation of this will follow in the future.

docs/core/index.qmd

@@ -5,7 +5,7 @@ title: "Julia core"
 With the term "core", we mean the computational engine of Ribasim. As detailed in the
 [usage](usage.qmd) documentation, it is generally used as a command line tool.
 
-The theory is described on the [equations](equations.qmd) page, and more in-depth numerical considerations are described on the [numerical considerations](numerics.qmd) page.
+The theory is described on the [equations](equations.qmd) page, and more in-depth numerical considerations are described on the [numerical considerations](numerics.qmd) page. As allocation is a large and self-contained part of the Ribasim core, it is described on the separate [allocation](allocation.qmd) page.
 
 The core is implemented in the [Julia programming language](https://julialang.org/), and
 can be found in the [Ribasim repository](https://github.com/Deltares/Ribasim) under the

python/ribasim_testmodels/ribasim_testmodels/allocation.py

@@ -535,11 +535,3 @@ def looped_subnetwork_model():
     )
 
     return model
-
-
-if __name__ == "__main__":
-    import matplotlib.pyplot as plt
-
-    model = looped_subnetwork_model()
-    model.plot()
-    plt.show()