Modeler allocation input (#699)

Fixes #656. --------- Co-authored-by: Martijn Visser <[email protected]>
Deltares · Oct 30, 2023 · c3c69a9 · c3c69a9
1 parent afe1bcc
commit c3c69a9
Show file tree

Hide file tree

Showing 23 changed files with 320 additions and 169 deletions.
diff --git a/build/libribasim/Manifest.toml b/build/libribasim/Manifest.toml
@@ -979,7 +979,7 @@ uuid = "ae029012-a4dd-5104-9daa-d747884805df"
 version = "1.3.0"
 
 [[deps.Ribasim]]
-deps = ["Arrow", "BasicModelInterface", "CodecLz4", "CodecZstd", "ComponentArrays", "Configurations", "DBInterface", "DataInterpolations", "DataStructures", "Dates", "Dictionaries", "DiffEqCallbacks", "FiniteDiff", "ForwardDiff", "Graphs", "HiGHS", "IterTools", "JuMP", "Legolas", "Logging", "LoggingExtras", "OrdinaryDiffEq", "PreallocationTools", "SQLite", "SciMLBase", "SparseArrays", "StructArrays", "Tables", "TerminalLoggers", "TimerOutputs", "TranscodingStreams"]
+deps = ["Arrow", "BasicModelInterface", "CodecLz4", "CodecZstd", "ComponentArrays", "Configurations", "DBInterface", "DataInterpolations", "DataStructures", "Dates", "Dictionaries", "DiffEqCallbacks", "FiniteDiff", "ForwardDiff", "Graphs", "HiGHS", "IterTools", "JuMP", "Legolas", "Logging", "LoggingExtras", "OrdinaryDiffEq", "PreallocationTools", "SQLite", "SciMLBase", "SparseArrays", "StructArrays", "Tables", "TerminalLoggers", "TimeZones", "TimerOutputs", "TranscodingStreams"]
 path = "../../core"
 uuid = "aac5e3d9-0b8f-4d4f-8241-b1a7a9632635"
 version = "0.3.0"

diff --git a/build/ribasim_cli/Manifest.toml b/build/ribasim_cli/Manifest.toml
@@ -979,7 +979,7 @@ uuid = "ae029012-a4dd-5104-9daa-d747884805df"
 version = "1.3.0"
 
 [[deps.Ribasim]]
-deps = ["Arrow", "BasicModelInterface", "CodecLz4", "CodecZstd", "ComponentArrays", "Configurations", "DBInterface", "DataInterpolations", "DataStructures", "Dates", "Dictionaries", "DiffEqCallbacks", "FiniteDiff", "ForwardDiff", "Graphs", "HiGHS", "IterTools", "JuMP", "Legolas", "Logging", "LoggingExtras", "OrdinaryDiffEq", "PreallocationTools", "SQLite", "SciMLBase", "SparseArrays", "StructArrays", "Tables", "TerminalLoggers", "TimerOutputs", "TranscodingStreams"]
+deps = ["Arrow", "BasicModelInterface", "CodecLz4", "CodecZstd", "ComponentArrays", "Configurations", "DBInterface", "DataInterpolations", "DataStructures", "Dates", "Dictionaries", "DiffEqCallbacks", "FiniteDiff", "ForwardDiff", "Graphs", "HiGHS", "IterTools", "JuMP", "Legolas", "Logging", "LoggingExtras", "OrdinaryDiffEq", "PreallocationTools", "SQLite", "SciMLBase", "SparseArrays", "StructArrays", "Tables", "TerminalLoggers", "TimeZones", "TimerOutputs", "TranscodingStreams"]
 path = "../../core"
 uuid = "aac5e3d9-0b8f-4d4f-8241-b1a7a9632635"
 version = "0.3.0"

diff --git a/core/Manifest.toml b/core/Manifest.toml
@@ -16,9 +16,9 @@ version = "0.4.4"
 
 [[deps.Adapt]]
 deps = ["LinearAlgebra", "Requires"]
-git-tree-sha1 = "76289dc51920fdc6e0013c872ba9551d54961c24"
+git-tree-sha1 = "02f731463748db57cc2ebfbd9fbc9ce8280d3433"
 uuid = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
-version = "3.6.2"
+version = "3.7.1"
 weakdeps = ["StaticArrays"]
 
     [deps.Adapt.extensions]
@@ -188,9 +188,9 @@ version = "1.0.5+0"
 
 [[deps.ComponentArrays]]
 deps = ["ArrayInterface", "ChainRulesCore", "ForwardDiff", "Functors", "LinearAlgebra", "PackageExtensionCompat", "StaticArrayInterface", "StaticArraysCore"]
-git-tree-sha1 = "3ee12062466a73b5f9c04113375e4934f174bc2c"
+git-tree-sha1 = "f1ac18d8d4e8a6303d6d2d1c870528c3c52f0895"
 uuid = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
-version = "0.15.3"
+version = "0.15.4"
 
     [deps.ComponentArrays.extensions]
     ComponentArraysAdaptExt = "Adapt"
@@ -300,9 +300,9 @@ version = "0.3.25"
 
 [[deps.DiffEqBase]]
 deps = ["ArrayInterface", "ChainRulesCore", "DataStructures", "DocStringExtensions", "EnumX", "EnzymeCore", "FastBroadcast", "ForwardDiff", "FunctionWrappers", "FunctionWrappersWrappers", "LinearAlgebra", "Logging", "Markdown", "MuladdMacro", "Parameters", "PreallocationTools", "PrecompileTools", "Printf", "RecursiveArrayTools", "Reexport", "Requires", "SciMLBase", "SciMLOperators", "Setfield", "SparseArrays", "Static", "StaticArraysCore", "Statistics", "Tricks", "TruncatedStacktraces", "ZygoteRules"]
-git-tree-sha1 = "95b6df71e218379a831874215b0effaac791d7d7"
+git-tree-sha1 = "36a590efdbee58b38f903ffc3b378f4a5336bc3f"
 uuid = "2b5f629d-d688-5b77-993f-72d75c75574e"
-version = "6.133.1"
+version = "6.134.0"
 
     [deps.DiffEqBase.extensions]
     DiffEqBaseDistributionsExt = "Distributions"
@@ -501,9 +501,9 @@ version = "1.9.0"
 
 [[deps.HiGHS]]
 deps = ["HiGHS_jll", "MathOptInterface", "PrecompileTools", "SparseArrays"]
-git-tree-sha1 = "9d75ef949c17a2a150b91b8365a6e5bc43a2a0d3"
+git-tree-sha1 = "fce13308f09771b160232903cad57be39a8a0ebb"
 uuid = "87dc4568-4c63-4d18-b0c0-bb2238e4078b"
-version = "1.7.3"
+version = "1.7.5"
 
 [[deps.HiGHS_jll]]
 deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "Libdl"]
@@ -577,9 +577,9 @@ version = "0.21.4"
 
 [[deps.JuMP]]
 deps = ["LinearAlgebra", "MacroTools", "MathOptInterface", "MutableArithmetics", "OrderedCollections", "Printf", "SnoopPrecompile", "SparseArrays"]
-git-tree-sha1 = "3700a700bc80856fe673b355123ae4574f2d5dfe"
+git-tree-sha1 = "25b2fcda4d455b6f93ac753730d741340ba4a4fe"
 uuid = "4076af6c-e467-56ae-b986-b466b2749572"
-version = "1.15.1"
+version = "1.16.0"
 
     [deps.JuMP.extensions]
     JuMPDimensionalDataExt = "DimensionalData"
@@ -661,9 +661,9 @@ uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 
 [[deps.LinearSolve]]
 deps = ["ArrayInterface", "ConcreteStructs", "DocStringExtensions", "EnumX", "EnzymeCore", "FastLapackInterface", "GPUArraysCore", "InteractiveUtils", "KLU", "Krylov", "Libdl", "LinearAlgebra", "MKL_jll", "PrecompileTools", "Preferences", "RecursiveFactorization", "Reexport", "Requires", "SciMLBase", "SciMLOperators", "Setfield", "SparseArrays", "Sparspak", "SuiteSparse", "UnPack"]
-git-tree-sha1 = "31353ba539d14a342908f765407abffd8db5f3c7"
+git-tree-sha1 = "a563cd835c9ed5295c35a7d140e0bf0a514bb10a"
 uuid = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
-version = "2.11.0"
+version = "2.13.0"
 
     [deps.LinearSolve.extensions]
     LinearSolveBandedMatricesExt = "BandedMatrices"
@@ -756,9 +756,9 @@ uuid = "d6f4376e-aef5-505a-96c1-9c027394607a"
 
 [[deps.MathOptInterface]]
 deps = ["BenchmarkTools", "CodecBzip2", "CodecZlib", "DataStructures", "ForwardDiff", "JSON", "LinearAlgebra", "MutableArithmetics", "NaNMath", "OrderedCollections", "PrecompileTools", "Printf", "SparseArrays", "SpecialFunctions", "Test", "Unicode"]
-git-tree-sha1 = "5c9f1e635e8d491297e596b56fec1c95eafb95a3"
+git-tree-sha1 = "13b3d40084d04e609e0509730f05215fb2a2fba4"
 uuid = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
-version = "1.20.1"
+version = "1.21.0"
 
 [[deps.MbedTLS_jll]]
 deps = ["Artifacts", "Libdl"]
@@ -813,9 +813,9 @@ version = "1.2.0"
 
 [[deps.NonlinearSolve]]
 deps = ["ADTypes", "ArrayInterface", "ConcreteStructs", "DiffEqBase", "EnumX", "FiniteDiff", "ForwardDiff", "LineSearches", "LinearAlgebra", "LinearSolve", "PrecompileTools", "RecursiveArrayTools", "Reexport", "SciMLBase", "SimpleNonlinearSolve", "SparseArrays", "SparseDiffTools", "StaticArraysCore", "UnPack"]
-git-tree-sha1 = "ee92770e0832314ccd424d83a0ab4c75fc6dc91f"
+git-tree-sha1 = "9203b3333c9610664de2e8cbc23cfd726663df7d"
 uuid = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
-version = "2.3.0"
+version = "2.4.0"
 
     [deps.NonlinearSolve.extensions]
     NonlinearSolveFastLevenbergMarquardtExt = "FastLevenbergMarquardt"
@@ -1021,17 +1021,19 @@ version = "3.43.0+0"
 
 [[deps.SciMLBase]]
 deps = ["ADTypes", "ArrayInterface", "ChainRulesCore", "CommonSolve", "ConstructionBase", "Distributed", "DocStringExtensions", "EnumX", "FillArrays", "FunctionWrappersWrappers", "IteratorInterfaceExtensions", "LinearAlgebra", "Logging", "Markdown", "PrecompileTools", "Preferences", "RecipesBase", "RecursiveArrayTools", "Reexport", "RuntimeGeneratedFunctions", "SciMLOperators", "StaticArraysCore", "Statistics", "SymbolicIndexingInterface", "Tables", "TruncatedStacktraces", "ZygoteRules"]
-git-tree-sha1 = "6134c8970f82f23c43d3580d79c3e47acf232083"
+git-tree-sha1 = "1c2a4e245744dd76b2eb677d3535ffad16d8b989"
 uuid = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
-version = "2.4.1"
+version = "2.5.0"
 
     [deps.SciMLBase.extensions]
+    SciMLBasePartialFunctionsExt = "PartialFunctions"
     SciMLBasePyCallExt = "PyCall"
     SciMLBasePythonCallExt = "PythonCall"
     SciMLBaseRCallExt = "RCall"
     SciMLBaseZygoteExt = "Zygote"
 
     [deps.SciMLBase.weakdeps]
+    PartialFunctions = "570af359-4316-4cb7-8c74-252c00c2016b"
     PyCall = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0"
     PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
     RCall = "6f49c342-dc21-5d91-9882-a32aef131414"
@@ -1076,9 +1078,9 @@ uuid = "1a1011a3-84de-559e-8e89-a11a2f7dc383"
 
 [[deps.SimpleNonlinearSolve]]
 deps = ["ArrayInterface", "DiffEqBase", "FiniteDiff", "ForwardDiff", "LinearAlgebra", "PackageExtensionCompat", "PrecompileTools", "Reexport", "SciMLBase", "StaticArraysCore"]
-git-tree-sha1 = "e308d089f5d0e733a017b61784c5813e672f760d"
+git-tree-sha1 = "15ff97fa4881133caa324dacafe28b5ac47ad8a2"
 uuid = "727e6d20-b764-4bd8-a329-72de5adea6c7"
-version = "0.1.22"
+version = "0.1.23"
 
     [deps.SimpleNonlinearSolve.extensions]
     SimpleNonlinearSolveNNlibExt = "NNlib"
@@ -1231,9 +1233,9 @@ version = "1.0.1"
 
 [[deps.Tables]]
 deps = ["DataAPI", "DataValueInterfaces", "IteratorInterfaceExtensions", "LinearAlgebra", "OrderedCollections", "TableTraits"]
-git-tree-sha1 = "a1f34829d5ac0ef499f6d84428bd6b4c71f02ead"
+git-tree-sha1 = "cb76cf677714c095e535e3501ac7954732aeea2d"
 uuid = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
-version = "1.11.0"
+version = "1.11.1"
 
 [[deps.Tar]]
 deps = ["ArgTools", "SHA"]
@@ -1338,9 +1340,9 @@ version = "1.5.5+0"
 
 [[deps.ZygoteRules]]
 deps = ["ChainRulesCore", "MacroTools"]
-git-tree-sha1 = "977aed5d006b840e2e40c0b48984f7463109046d"
+git-tree-sha1 = "9d749cd449fb448aeca4feee9a2f4186dbb5d184"
 uuid = "700de1a5-db45-46bc-99cf-38207098b444"
-version = "0.2.3"
+version = "0.2.4"
 
 [[deps.libblastrampoline_jll]]
 deps = ["Artifacts", "Libdl"]

diff --git a/core/src/Ribasim.jl b/core/src/Ribasim.jl
@@ -17,6 +17,7 @@ module Ribasim
 import IterTools
 import BasicModelInterface as BMI
 import HiGHS
+import JuMP
 import JuMP.Model as JuMPModel
 import TranscodingStreams
 
@@ -43,7 +44,6 @@ using Graphs:
     outneighbors,
     rem_edge!
 
-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

diff --git a/core/src/allocation.jl b/core/src/allocation.jl
@@ -382,7 +382,7 @@ function add_variables_flow!(
     allocgraph_edges::Vector{Edge{Int}},
 )::Nothing
     n_flows = length(allocgraph_edges)
-    problem[:F] = @variable(problem, F[1:n_flows] >= 0.0)
+    problem[:F] = JuMP.@variable(problem, F[1:n_flows] >= 0.0)
     return nothing
 end
 
@@ -401,7 +401,7 @@ function add_variables_allocation_user!(
         allocgraph_node_id_user = allocgraph_edges[allocgraph_edge_id_user_demand].dst
         base_name = "A_user_$allocgraph_node_id_user"
         problem[Symbol(base_name)] =
-            @variable(problem, [1:length(user.priorities)], base_name = base_name)
+            JuMP.@variable(problem, [1:length(user.priorities)], base_name = base_name)
     end
     return nothing
 end
@@ -416,7 +416,7 @@ function add_variables_allocation_basin!(
     node_id_mapping::Dict{Int, Tuple{Int, Symbol}},
     allocgraph_node_ids_basin::Vector{Int},
 )::Nothing
-    @variable(problem, A_basin[i = allocgraph_node_ids_basin] >= 0.0)
+    JuMP.@variable(problem, A_basin[i = allocgraph_node_ids_basin] >= 0.0)
     return nothing
 end
 
@@ -447,16 +447,16 @@ function add_constraints_user_allocation!(
         base_name = "A_user_$allocgraph_node_id_user"
         A_user = problem[Symbol(base_name)]
         # Sum of allocations to user is total flow to user
-        @constraint(
+        JuMP.@constraint(
             problem,
             sum(A_user) == F[allocgraph_edge_id_user_demand],
             base_name = "allocation_sum[$allocgraph_node_id_user]"
         )
         # Allocation flows are non-negative
-        @constraint(problem, [p = 1:length(user.priorities)], A_user[p] >= 0)
+        JuMP.@constraint(problem, [p = 1:length(user.priorities)], A_user[p] >= 0)
         # Allocation flows are bounded from above by demands
         base_name = "demand_user_$allocgraph_node_id_user"
-        problem[Symbol(base_name)] = @constraint(
+        problem[Symbol(base_name)] = JuMP.@constraint(
             problem,
             [p = 1:length(user.priorities)],
             A_user[p] <= 0,
@@ -480,7 +480,7 @@ function add_constraints_basin_allocation!(
     allocgraph_node_ids_basin::Vector{Int},
 )::Nothing
     A_basin = problem[:A_basin]
-    problem[:basin_allocation] = @constraint(
+    problem[:basin_allocation] = JuMP.@constraint(
         problem,
         [i = allocgraph_node_ids_basin],
         A_basin[i] <= 0.0,
@@ -509,7 +509,7 @@ function add_constraints_capacity!(
             push!(allocgraph_edge_ids_finite_capacity, i)
         end
     end
-    problem[:capacity] = @constraint(
+    problem[:capacity] = JuMP.@constraint(
         problem,
         [i = allocgraph_edge_ids_finite_capacity],
         F[i] <= capacity[allocgraph_edges[i].src, allocgraph_edges[i].dst],
@@ -533,7 +533,7 @@ function add_constraints_source!(
     graph_allocation::DiGraph{Int},
 )::Nothing
     F = problem[:F]
-    problem[:source] = @constraint(
+    problem[:source] = JuMP.@constraint(
         problem,
         [i = keys(source_edge_mapping)],
         F[findfirst(
@@ -559,7 +559,7 @@ function add_constraints_flow_conservation!(
     allocgraph_node_outedge_ids::Dict{Int, Vector{Int}},
 )::Nothing
     F = problem[:F]
-    problem[:flow_conservation] = @constraint(
+    problem[:flow_conservation] = JuMP.@constraint(
         problem,
         [i = allocgraph_node_ids_basin],
         sum([
@@ -585,7 +585,7 @@ function add_constraints_user_returnflow!(
 )::Nothing
     F = problem[:F]
 
-    problem[:return_flow] = @constraint(
+    problem[:return_flow] = JuMP.@constraint(
         problem,
         [i = allocgraph_node_ids_user_with_returnflow],
         F[only(allocgraph_node_outedge_ids[i])] ==
@@ -612,7 +612,7 @@ function add_objective_function!(
         problem[Symbol("A_user_$allocgraph_node_id_user")] for
         allocgraph_node_id_user in allocgraph_node_ids_user
     ]
-    @objective(
+    JuMP.@objective(
         problem,
         Max,
         sum(A_basin) + sum([
@@ -647,7 +647,8 @@ function allocation_problem(
     allocgraph_node_inedge_ids, allocgraph_node_outedge_ids =
         get_node_in_out_edges(graph_allocation)
 
-    problem = JuMPModel(HiGHS.Optimizer)
+    optimizer = JuMP.optimizer_with_attributes(HiGHS.Optimizer, "log_to_console" => false)
+    problem = JuMPModel(optimizer)
 
     # Add variables to problem
     add_variables_flow!(problem, allocgraph_edges)
@@ -778,7 +779,7 @@ function set_model_state_in_allocation!(
             base_name = "demand_user_$allocgraph_node_id"
             constraints_demand = problem[Symbol(base_name)]
             for priority_idx in eachindex(priorities)
-                set_normalized_rhs(
+                JuMP.set_normalized_rhs(
                     constraints_demand[priority_idx],
                     demand_node[priority_idx](t),
                 )
@@ -787,7 +788,7 @@ function set_model_state_in_allocation!(
             subnetwork_edge = source_edge_mapping[allocgraph_node_id]
             subnetwork_node_ids = edge_ids_flow_inv[subnetwork_edge]
             constraint_source = problem[:source][allocgraph_node_id]
-            set_normalized_rhs(constraint_source, flow[subnetwork_node_ids])
+            JuMP.set_normalized_rhs(constraint_source, flow[subnetwork_node_ids])
         elseif allocgraph_node_type == :basin
             # TODO: Compute basin flow from vertical fluxes and basin volume.
             # Set as basin demand if the net flow is negative, set as source
@@ -810,7 +811,7 @@ function assign_allocations!(allocation_model::AllocationModel, user::User)::Not
         if allocgraph_node_type == :user
             user_idx = findsorted(user.node_id, subnetwork_node_id)
             base_name = "A_user_$allocgraph_node_id"
-            user.allocated[user_idx] .= value.(problem[Symbol(base_name)])
+            user.allocated[user_idx] .= JuMP.value.(problem[Symbol(base_name)])
         end
     end
     return nothing
@@ -821,11 +822,17 @@ Update the allocation optimization problem for the given subnetwork with the pro
 and flows, solve the allocation problem and assign the results to the users.
 """
 function allocate!(p::Parameters, allocation_model::AllocationModel, t::Float64)::Nothing
+    (; problem) = allocation_model
+
     # Update allocation problem with data from main model
     set_model_state_in_allocation!(allocation_model, p, t)
 
     # Solve the allocation problem
-    optimize!(allocation_model.problem)
+    JuMP.optimize!(problem)
+    @debug JuMP.solution_summary(problem)
+    if JuMP.termination_status(problem) !== JuMP.OPTIMAL
+        error("Allocation coudn't find optimal solution.")
+    end
 
     # Assign the allocations to the users
     assign_allocations!(allocation_model, p.user)