diff --git a/core/src/allocation_optim.jl b/core/src/allocation_optim.jl index d0d23a2e1..7ed7ac55c 100644 --- a/core/src/allocation_optim.jl +++ b/core/src/allocation_optim.jl @@ -524,13 +524,14 @@ Set the demand of the flow demand nodes. 2 cases: - Before an allocation solve, subtract the flow trough the node with a flow demand from the total flow demand (which will be used at the priority of the flow demand only). """ -function adjust_demands_user!( +function adjust_demands!( allocation_model::AllocationModel, p::Parameters, priority_idx::Int, + user_demand::UserDemand, )::Nothing (; problem, subnetwork_id) = allocation_model - (; graph, user_demand) = p + (; graph) = p (; node_id, demand_reduced) = user_demand F = problem[:F] @@ -551,7 +552,13 @@ end Subtract the allocated flow to the basin from its demand, to obtain the reduced demand used for goal programming """ -function adjust_demands_level!(allocation_model::AllocationModel, p::Parameters)::Nothing + +function adjust_demands!( + allocation_model::AllocationModel, + p::Parameters, + ::Int, + ::LevelDemand, +)::Nothing (; graph, basin) = p (; node_id, demand) = basin (; subnetwork_id, problem) = allocation_model @@ -593,8 +600,13 @@ end Reduce the flow demand based on flow trough the node with the demand. Flow from any priority counts. """ -function adjust_demands_flow!(allocation_model::AllocationModel, p::Parameters)::Nothing - (; flow_demand, graph) = p +function adjust_demands!( + allocation_model::AllocationModel, + p::Parameters, + ::Int, + flow_demand::FlowDemand, +)::Nothing + (; graph) = p (; problem, subnetwork_id) = allocation_model F = problem[:F] @@ -861,7 +873,7 @@ function save_allocation_flows!( return nothing end -function allocate_priority!( +function optimize_priority!( allocation_model::AllocationModel, u::ComponentVector, p::Parameters, @@ -916,9 +928,12 @@ function allocate_priority!( adjust_capacities_returnflow!(allocation_model, p) # Adjust demands for next optimization (in case of internal_sources -> collect_demands) - adjust_demands_user!(allocation_model, p, priority_idx) - adjust_demands_level!(allocation_model, p) - adjust_demands_flow!(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 return nothing end @@ -968,43 +983,71 @@ end Update the allocation optimization problem for the given subnetwork with the problem state and flows, solve the allocation problem and assign the results to the UserDemand. """ -function allocate!( +function collect_demands!( p::Parameters, allocation_model::AllocationModel, t::Float64, u::ComponentVector, - optimization_type::OptimizationType.T, )::Nothing (; allocation) = p (; subnetwork_id) = allocation_model (; priorities, subnetwork_demands) = allocation - main_network_source_edges = get_main_network_connections(p, subnetwork_id) - if subnetwork_id == 1 - @assert optimization_type == OptimizationType.allocate "For the main network no demands have to be collected" + ## Find internal sources + optimization_type = OptimizationType.internal_sources + set_initial_capacities_inlet!(allocation_model, p, optimization_type) + set_initial_values!(allocation_model, p, u, t) + + # Loop over priorities + for priority_idx in eachindex(priorities) + optimize_priority!(allocation_model, u, p, t, priority_idx, optimization_type) end + ## Collect demand + optimization_type = OptimizationType.collect_demands + + main_network_source_edges = get_main_network_connections(p, subnetwork_id) + # Reset the subnetwork demands to 0.0 - if optimization_type == OptimizationType.collect_demands - for main_network_connection in keys(subnetwork_demands) - if main_network_connection in main_network_source_edges - subnetwork_demands[main_network_connection] .= 0.0 - end + for main_network_connection in keys(subnetwork_demands) + if main_network_connection in main_network_source_edges + subnetwork_demands[main_network_connection] .= 0.0 end end set_initial_capacities_inlet!(allocation_model, p, optimization_type) - if optimization_type == OptimizationType.collect_demands - # When collecting demands, only flow should be available - # from the main to subnetwork connections - empty_sources!(allocation_model, allocation) - else - set_initial_values!(allocation_model, p, u, t) + # When collecting demands, only flow should be available + # from the main to subnetwork connections + empty_sources!(allocation_model, allocation) + + # Loop over priorities + for priority_idx in eachindex(priorities) + optimize_priority!(allocation_model, u, p, t, priority_idx, optimization_type) + end +end + +function allocate_demands!( + p::Parameters, + allocation_model::AllocationModel, + t::Float64, + u::ComponentVector, +)::Nothing + optimization_type = OptimizationType.allocate + (; allocation) = p + (; subnetwork_id) = allocation_model + (; priorities) = allocation + + if is_main_network(subnetwork_id) + @assert optimization_type == OptimizationType.allocate "For the main network no demands have to be collected" end + set_initial_capacities_inlet!(allocation_model, p, optimization_type) + + set_initial_values!(allocation_model, p, u, t) + # Loop over the priorities for priority_idx in eachindex(priorities) - allocate_priority!(allocation_model, u, p, t, priority_idx, optimization_type) + optimize_priority!(allocation_model, u, p, t, priority_idx, optimization_type) end end diff --git a/core/src/callback.jl b/core/src/callback.jl index 578b73243..953b2248f 100644 --- a/core/src/callback.jl +++ b/core/src/callback.jl @@ -480,8 +480,7 @@ function update_allocation!(integrator)::Nothing # If a main network is present, collect demands of subnetworks if has_main_network(allocation) for allocation_model in Iterators.drop(allocation_models, 1) - allocate!(p, allocation_model, t, u, OptimizationType.internal_sources) - allocate!(p, allocation_model, t, u, OptimizationType.collect_demands) + collect_demands!(p, allocation_model, t, u) end end @@ -489,7 +488,7 @@ function update_allocation!(integrator)::Nothing # If a main network is present this is solved first, # which provides allocation to the subnetworks for allocation_model in allocation_models - allocate!(p, allocation_model, t, u, OptimizationType.allocate) + allocate_demands!(p, allocation_model, t, u) end # Reset the mean source flows diff --git a/core/src/parameter.jl b/core/src/parameter.jl index 8108336ce..b280625f7 100644 --- a/core/src/parameter.jl +++ b/core/src/parameter.jl @@ -135,6 +135,8 @@ end abstract type AbstractParameterNode end +abstract type AbstractDemandNode <: AbstractParameterNode end + """ In-memory storage of saved mean flows for writing to results. @@ -520,7 +522,7 @@ allocated: water flux currently allocated to UserDemand per priority (node_idx, return_factor: the factor in [0,1] of how much of the abstracted water is given back to the system min_level: The level of the source basin below which the UserDemand does not abstract """ -struct UserDemand <: AbstractParameterNode +struct UserDemand <: AbstractDemandNode node_id::Vector{NodeID} inflow_edge::Vector{EdgeMetadata} outflow_edge::Vector{EdgeMetadata} @@ -576,14 +578,14 @@ min_level: The minimum target level of the connected basin(s) max_level: The maximum target level of the connected basin(s) priority: If in a shortage state, the priority of the demand of the connected basin(s) """ -struct LevelDemand <: AbstractParameterNode +struct LevelDemand <: AbstractDemandNode node_id::Vector{NodeID} min_level::Vector{ScalarInterpolation} max_level::Vector{ScalarInterpolation} priority::Vector{Int32} end -struct FlowDemand <: AbstractParameterNode +struct FlowDemand <: AbstractDemandNode node_id::Vector{NodeID} demand_itp::Vector{ScalarInterpolation} demand::Vector{Float64} diff --git a/core/test/allocation_test.jl b/core/test/allocation_test.jl index a3ef4f1af..5eb6a8a87 100644 --- a/core/test/allocation_test.jl +++ b/core/test/allocation_test.jl @@ -17,7 +17,7 @@ allocation.mean_flows[(NodeID(:FlowBoundary, 1), NodeID(:Basin, 2))][] = 4.5 allocation_model = p.allocation.allocation_models[1] u = ComponentVector(; storage = zeros(length(p.basin.node_id))) - Ribasim.allocate!(p, allocation_model, 0.0, u, OptimizationType.allocate) + Ribasim.allocate_demands!(p, allocation_model, 0.0, u) # Last priority (= 2) flows F = allocation_model.problem[:F] @@ -223,8 +223,7 @@ end # Collecting demands u = ComponentVector(; storage = zeros(length(basin.node_id))) for allocation_model in allocation_models[2:end] - Ribasim.allocate!(p, allocation_model, t, u, OptimizationType.internal_sources) - Ribasim.allocate!(p, allocation_model, t, u, OptimizationType.collect_demands) + Ribasim.collect_demands!(p, allocation_model, t, u) end # See the difference between these values here and in @@ -238,7 +237,7 @@ end # Solving for the main network, containing subnetworks as UserDemands allocation_model = allocation_models[1] (; problem) = allocation_model - Ribasim.allocate_priority!(allocation_model, u, p, t, 1, OptimizationType.allocate) + Ribasim.optimize_priority!(allocation_model, u, p, t, 1, OptimizationType.allocate) # Main network objective function F = problem[:F] @@ -307,8 +306,7 @@ end # Collecting demands u = ComponentVector(; storage = zeros(length(basin.node_id))) for allocation_model in allocation_models[2:end] - Ribasim.allocate!(p, allocation_model, t, u, OptimizationType.internal_sources) - Ribasim.allocate!(p, allocation_model, t, u, OptimizationType.collect_demands) + Ribasim.collect_demands!(p, allocation_model, t, u) end # See the difference between these values here and in @@ -453,7 +451,7 @@ end Ribasim.set_initial_values!(allocation_model, p, u, t) # Priority 1 - Ribasim.allocate_priority!( + Ribasim.optimize_priority!( allocation_model, model.integrator.u, p, @@ -468,7 +466,7 @@ end @test JuMP.normalized_rhs(constraint_flow_out) == Inf ## Priority 2 - Ribasim.allocate_priority!( + Ribasim.optimize_priority!( allocation_model, model.integrator.u, p, @@ -483,7 +481,7 @@ end @test JuMP.normalized_rhs(constraint_flow_out) == 0.0 ## Priority 3 - Ribasim.allocate_priority!( + Ribasim.optimize_priority!( allocation_model, model.integrator.u, p, @@ -502,7 +500,7 @@ end @test JuMP.value(F[(only(level_boundary.node_id), node_id_with_flow_demand)]) == 0 ## Priority 4 - Ribasim.allocate_priority!( + Ribasim.optimize_priority!( allocation_model, model.integrator.u, p, diff --git a/docs/contribute/addnode.qmd b/docs/contribute/addnode.qmd index 1bd65f434..bf62d262e 100644 --- a/docs/contribute/addnode.qmd +++ b/docs/contribute/addnode.qmd @@ -16,6 +16,14 @@ struct NewNodeType <: AbstractParameterNode # Other fields end ``` +Another abstract type which subtypes from `AbstractParameterNode` is called `AbstractDemandNode`. For creating new node type used in allocation, define a struct: + +```julia +struct NewNodeType <: AbstractDemandNode + node_id::Vector{NodeID} + # Other fields +end +``` These fields do not have to correspond 1:1 with the input tables (see below). The vector with all node IDs that are of the new type in a given model is a mandatory field. Now you can: diff --git a/pixi.lock b/pixi.lock index 762397b90..a11d26200 100644 --- a/pixi.lock +++ b/pixi.lock @@ -8458,28 +8458,6 @@ packages: license_family: BSD size: 15322 timestamp: 1712542844430 -- kind: conda - name: bleach - version: 6.1.0 - build: pyhd8ed1ab_0 - subdir: noarch - noarch: python - url: https://conda.anaconda.org/conda-forge/noarch/bleach-6.1.0-pyhd8ed1ab_0.conda - sha256: 845e77ef495376c5c3c328ccfd746ca0ef1978150cae8eae61a300fe7755fb08 - md5: 0ed9d7c0e9afa7c025807a9a8136ea3e - depends: - - packaging - - python >=3.6 - - setuptools - - six >=1.9.0 - - webencodings - license: Apache-2.0 - license_family: Apache - purls: - - pkg:pypi/html5lib - - pkg:pypi/bleach - size: 131220 - timestamp: 1696630354218 - kind: conda name: bleach version: 6.1.0 @@ -33394,11 +33372,11 @@ packages: - matplotlib - numpy - pandas - - pandera>=0.18 + - pandera >=0.18 - pyarrow - - pydantic~=2.0 + - pydantic ~=2.0 - pyogrio - - shapely>=2.0 + - shapely >=2.0 - tomli - tomli-w - pytest ; extra == 'tests'