Modeler allocation input

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"]

core/src/Ribasim.jl

@@ -43,7 +43,17 @@ using Graphs:
     outneighbors,
     rem_edge!
 
-using JuMP: @variable, @constraint, @objective, set_normalized_rhs, optimize!, value
+using JuMP:
+    @variable,
+    @constraint,
+    @objective,
+    set_normalized_rhs,
+    optimize!,
+    value,
+    OPTIMAL,
+    optimizer_with_attributes,
+    solution_summary,
+    termination_status
 using Legolas: Legolas, @schema, @version, validate, SchemaVersion, declared
 using Logging: current_logger, min_enabled_level, with_logger
 using LoggingExtras: EarlyFilteredLogger, LevelOverrideLogger

core/src/allocation.jl

@@ -647,7 +647,8 @@
     allocgraph_node_inedge_ids, allocgraph_node_outedge_ids =
         get_node_in_out_edges(graph_allocation)
 
-    problem = JuMPModel(HiGHS.Optimizer)
+    optimizer = optimizer_with_attributes(HiGHS.Optimizer, "log_to_console" => false)
+    problem = JuMPModel(optimizer)
 
     # Add variables to problem
     add_variables_flow!(problem, allocgraph_edges)
@@ -821,11 +822,17 @@
 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)
+    optimize!(problem)
+    @debug solution_summary(problem)
+    if termination_status(problem) !== OPTIMAL
+        error("Allocation coudn't find optimal solution.")
+    end
 
     # Assign the allocations to the users
     assign_allocations!(allocation_model, p.user)

core/src/bmi.jl

@@ -82,6 +82,11 @@
         # use state
         state = load_structvector(db, config, BasinStateV1)
         n = length(get_ids(db, "Basin"))
+
+        # Allocation data structures
+        if config.allocation.use_allocation
+            generate_allocation_models!(parameters, db, config)
+        end
     finally
         # always close the database, also in case of an error
         close(db)
@@ -116,7 +121,7 @@
     end
     @debug "Setup ODEProblem."
 
-    callback, saved_flow = create_callbacks(parameters; config.solver.saveat)
+    callback, saved_flow = create_callbacks(parameters, config; config.solver.saveat)
     @debug "Created callbacks."
 
     # Initialize the integrator, providing all solver options as described in
@@ -209,7 +214,8 @@
 Returns the CallbackSet and the SavedValues for flow.
 """
 function create_callbacks(
-    parameters;
+    parameters,
+    config;
     saveat,
 )::Tuple{CallbackSet, SavedValues{Float64, Vector{Float64}}}
     (; starttime, basin, tabulated_rating_curve, discrete_control) = parameters
@@ -223,6 +229,13 @@
     tabulated_rating_curve_cb = PresetTimeCallback(tstops, update_tabulated_rating_curve!)
     push!(callbacks, tabulated_rating_curve_cb)
 
+    if config.allocation.use_allocation
+        t_end = seconds_since(config.endtime, config.starttime)
+        tstops = 1:(config.allocation.timestep):t_end
+        allocation_cb = PresetTimeCallback(tstops, update_allocation!)
+        push!(callbacks, allocation_cb)
+    end
+
     # save the flows over time, as a Vector of the nonzeros(flow)
     saved_flow = SavedValues(Float64, Vector{Float64})
     save_flow_cb = SavingCallback(save_flow, saved_flow; saveat, save_start = false)
@@ -502,6 +515,14 @@
     return nothing
 end
 
+"Solve the allocation problem for all users and assign allocated abstractions to user nodes."
+function update_allocation!(integrator)::Nothing
+    (; p, t) = integrator
+    for allocation_model in integrator.p.connectivity.allocation_models
+        allocate!(p, allocation_model, t)
+    end
+end
+
 "Load updates from 'TabulatedRatingCurve / time' into the parameters"
 function update_tabulated_rating_curve!(integrator)::Nothing
     (; node_id, tables, time) = integrator.p.tabulated_rating_curve

core/src/config.jl

@@ -122,13 +122,15 @@ end
     timing::Bool = false
 end
 
+@option struct Allocation <: TableOption
+    timestep::Union{Float64, Nothing} = nothing
+    use_allocation::Bool = false
+end
+
 @option @addnodetypes struct Config <: TableOption
     starttime::DateTime
     endtime::DateTime
 
-    # [s] Δt for periodic update frequency, including user horizons
-    update_timestep::Float64 = 60 * 60 * 24.0
-
     # optional, when Config is created from a TOML file, this is its directory
     relative_dir::String = "."  # ignored(!)
     input_dir::String = "."
@@ -137,12 +139,12 @@ end
     # input, required
     database::String
 
-    # results, required
-    results::Results = Results()
-
+    allocation::Allocation = Allocation()
     solver::Solver = Solver()
-
     logging::Logging = Logging()
+
+    # results, required
+    results::Results = Results()
 end
 
 function Configurations.from_dict(::Type{Logging}, ::Type{LogLevel}, level::AbstractString)

core/src/create.jl

@@ -230,7 +230,6 @@
         flow = FixedSizeDiffCache(flow, chunk_size)
     end
 
-    # TODO: Create allocation models from input here
     allocation_models = AllocationModel[]
 
     return Connectivity(
@@ -246,6 +245,53 @@
     )
 end
 
+function generate_allocation_models!(p::Parameters, db::DB, config::Config)::Nothing
+    (; connectivity) = p
+    node = load_structvector(db, config, NodeV1)
+    edge = load_structvector(db, config, EdgeV1)
+
+    # coalesce control_state to nothing to avoid boolean groupby logic on missing
+    allocation_groups_node =
+        IterTools.groupby(row -> coalesce(row.allocation_network_id, nothing), node)
+    allocation_groups_edge =
+        IterTools.groupby(row -> coalesce(row.allocation_network_id, nothing), edge)
+
+    allocation_groups_node_dict = Dict{Int, StructVector}()
+    allocation_groups_edge_dict = Dict{Int, StructVector}()
+
+    for allocation_group_node in allocation_groups_node
+        allocation_network_id = first(allocation_group_node).allocation_network_id
+        if !ismissing(allocation_network_id)
+            allocation_groups_node_dict[allocation_network_id] = allocation_group_node
+        end
+    end
+    for allocation_group_edge in allocation_groups_edge
+        allocation_network_id = first(allocation_group_edge).allocation_network_id
+        if !ismissing(allocation_network_id)
+            allocation_groups_edge_dict[allocation_network_id] = allocation_group_edge
+        end
+    end
+
+    for (allocation_network_id, allocation_group_node) in allocation_groups_node_dict
+        allocation_group_edge = get(
+            allocation_groups_edge_dict,
+            allocation_network_id,
+            StructVector{EdgeV1}(undef, 0),
+        )
+        source_edge_ids = [row.fid for row in allocation_group_edge]
+        push!(
+            connectivity.allocation_models,
+            AllocationModel(
+                p,
+                allocation_group_node.fid,
+                source_edge_ids,
+                config.allocation.timestep,
+            ),
+        )
+    end
+    return nothing
+end
+
 function LinearResistance(db::DB, config::Config)::LinearResistance
     static = load_structvector(db, config, LinearResistanceStaticV1)
     parsed_parameters, valid = parse_static_and_time(db, config, "LinearResistance"; static)

core/src/validation.jl

@@ -137,11 +137,11 @@ n_neighbor_bounds_control(::Val{:User}) = n_neighbor_bounds(0, 0, 0, 0)
 n_neighbor_bounds_control(nodetype) =
     error("'n_neighbor_bounds_control' not defined for $nodetype.")
 
-# TODO NodeV1 and EdgeV1 are not yet used
 @version NodeV1 begin
     fid::Int
     name::String = isnothing(s) ? "" : String(s)
     type::String = in(Symbol(type), nodetypes) ? type : error("Unknown node type $type")
+    allocation_network_id::Union{Missing, Int}
 end
 
 @version EdgeV1 begin
@@ -150,6 +150,7 @@ end
     from_node_id::Int
     to_node_id::Int
     edge_type::String
+    allocation_network_id::Union{Missing, Int}
 end
 
 @version PumpStaticV1 begin

core/test/allocation.jl

@@ -25,10 +25,10 @@ using JuMP: value
     Ribasim.allocate!(p, allocation_model, t)
 
     F = value.(allocation_model.problem[:F])
-    @test F ≈ [3.0, 3.0, 0.5, 3.5, 1.0, 4.5]
+    @test F ≈ [-0.0, 0.0, 4.0, 4.0, 0.5, 4.5]
 
     allocated = p.user.allocated
-    @test allocated[1] ≈ [0.0, 1.0, 0.0]
-    @test allocated[2] ≈ [0.0, 0.5, 0.0]
-    @test allocated[3] ≈ [3.0, 0.0, 0.0]
+    @test allocated[1] ≈ [0.0, 0.5]
+    @test allocated[2] ≈ [4.0, 0.0]
+    @test allocated[3] ≈ [0.0, 0.0]
 end