Merge branch 'main' into not-so-auto-update

core/src/allocation_optim.jl

@@ -717,11 +717,9 @@ function save_allocation_flows!(
     priority::Int32,
     optimization_type::OptimizationType.T,
 )::Nothing
-    (; flow, problem, subnetwork_id) = allocation_model
+    (; flow, subnetwork_id, sources) = allocation_model
     (; allocation, graph) = p
     (; record_flow) = allocation
-    F_basin_in = problem[:F_basin_in]
-    F_basin_out = problem[:F_basin_out]
 
     edges_allocation = keys(flow.data)
 
@@ -776,7 +774,7 @@ function save_allocation_flows!(
     for node_id in graph[].node_ids[subnetwork_id]
         if node_id.type == NodeType.Basin &&
            has_external_demand(graph, node_id, :level_demand)[1]
-            flow_rate = JuMP.value(F_basin_out[node_id]) - JuMP.value(F_basin_in[node_id])
+            flow_rate = sources[(node_id, node_id)].basin_flow_rate
             push!(record_flow.time, t)
             push!(record_flow.edge_id, 0)
             push!(record_flow.from_node_type, string(NodeType.Basin))
@@ -894,6 +892,15 @@ function optimize_per_source!(
 )::Nothing
     (; problem, sources, subnetwork_id, flow) = allocation_model
     (; priorities) = allocation
+    F_basin_in = problem[:F_basin_in]
+    F_basin_out = problem[:F_basin_out]
+
+    # Start the cumulative basin flow rates at 0
+    for source in values(sources)
+        if source.type == AllocationSourceType.basin
+            source.basin_flow_rate = 0.0
+        end
+    end
 
     priority = priorities[priority_idx]
 
@@ -940,6 +947,15 @@ function optimize_per_source!(
             end
         end
 
+        # Add to the basin cumulative flow rate
+        for (edge, source) in sources
+            if source.type == AllocationSourceType.basin
+                node_id = edge[1]
+                source.basin_flow_rate +=
+                    JuMP.value(F_basin_out[node_id]) - JuMP.value(F_basin_in[node_id])
+            end
+        end
+
         # Adjust allocated flow to basins
         increase_allocateds!(p.basin, problem)
     end

core/src/parameter.jl

@@ -136,12 +136,15 @@ edge: The outflow edge of the source
 type: The type of source (edge, basin, main_to_sub, user_return, buffer)
 capacity: The initial capacity of the source as determined by the physical layer
 capacity_reduced: The capacity adjusted by passed optimizations
+basin_flow_rate: The total outflow rate of a basin when optimized over all sources for one priority.
+    Ignored when the source is not a basin.
 """
 @kwdef mutable struct AllocationSource
     const edge::Tuple{NodeID, NodeID}
     const type::AllocationSourceType.T
     capacity::Float64 = 0.0
     capacity_reduced::Float64 = 0.0
+    basin_flow_rate::Float64 = 0.0
 end
 
 function Base.show(io::IO, source::AllocationSource)

python/ribasim_testmodels/ribasim_testmodels/__init__.py

@@ -7,6 +7,7 @@
 import ribasim_testmodels
 from ribasim_testmodels.allocation import (
     allocation_example_model,
+    allocation_training_model,
     fair_distribution_model,
     flow_demand_model,
     level_demand_model,
@@ -63,6 +64,7 @@
 from ribasim_testmodels.two_basin import two_basin_model
 
 __all__ = [
+    "allocation_training_model",
     "allocation_example_model",
     "backwater_model",
     "basic_arrow_model",

python/ribasim_testmodels/ribasim_testmodels/allocation.py

@@ -980,3 +980,181 @@ def fair_distribution_model():
     model.edge.add(model.user_demand[9], model.basin[5])
 
     return model
+
+
+def allocation_training_model():
+    model = Model(
+        starttime="2022-01-01",
+        endtime="2023-01-01",
+        crs="EPSG:4326",
+        allocation=Allocation(use_allocation=True),
+    )
+
+    flow_boundary_times = pd.date_range(start="2022-01-01", end="2023-01-01", freq="MS")
+
+    # Flow boundaries
+    main = model.flow_boundary.add(
+        Node(1, Point(0.0, 0.0), subnetwork_id=1, name="Main"),
+        [
+            flow_boundary.Time(
+                time=flow_boundary_times,
+                flow_rate=[
+                    47.3,
+                    156.7,
+                    77.6,
+                    47.8,
+                    26.6,
+                    23.1,
+                    18.6,
+                    15.6,
+                    23.1,
+                    35.6,
+                    24.4,
+                    20.0,
+                    29.4,
+                ],
+            )
+        ],
+    )
+
+    minor = model.flow_boundary.add(
+        Node(2, Point(-3.0, 0.0), subnetwork_id=1, name="Minor"),
+        [
+            flow_boundary.Time(
+                time=flow_boundary_times,
+                flow_rate=[
+                    0.2,
+                    28.3,
+                    16.0,
+                    11.2,
+                    8.5,
+                    9.6,
+                    9.2,
+                    7.9,
+                    7.5,
+                    7.2,
+                    7.4,
+                    10.0,
+                    8.3,
+                ],
+            )
+        ],
+    )
+
+    level = model.level_demand.add(
+        Node(11, Point(1, 1), subnetwork_id=1, name="test"),
+        [
+            level_demand.Static(
+                min_level=[2],
+                max_level=5,
+                priority=1,
+            )
+        ],
+    )
+
+    # Confluence
+    conf = model.basin.add(
+        Node(3, Point(-1.5, -1), subnetwork_id=1, name="confluence"),
+        [
+            basin.Profile(area=[672000, 5600000], level=[0, 6]),
+            basin.State(level=[4]),
+        ],
+    )
+
+    tbr_conf = model.tabulated_rating_curve.add(
+        Node(4, Point(-1.5, -1.5), subnetwork_id=1, name="tbr_conf"),
+        [
+            tabulated_rating_curve.Static(
+                level=[0.0, 2, 5],
+                flow_rate=[0.0, 50, 200],
+            )
+        ],
+    )
+
+    # Irrigation
+    irr = model.user_demand.add(
+        Node(6, Point(-1.5, 0.5), subnetwork_id=1, name="irrigation"),
+        [
+            user_demand.Time(
+                demand=[0.0, 0.0, 10, 12, 12, 0.0],
+                return_factor=0,
+                min_level=0,
+                priority=3,
+                time=[
+                    "2022-01-01",
+                    "2022-03-31",
+                    "2022-04-01",
+                    "2022-07-01",
+                    "2022-09-30",
+                    "2022-10-01",
+                ],
+            )
+        ],
+    )
+
+    # Reservoir
+    reservoir = model.basin.add(
+        Node(7, Point(-0.75, -0.5), subnetwork_id=1, name="reservoir"),
+        [
+            basin.Profile(area=[20000000, 32300000], level=[0, 7]),
+            basin.State(level=[3.5]),
+        ],
+    )
+
+    rsv_weir = model.tabulated_rating_curve.add(
+        Node(8, Point(-1.125, -0.75), subnetwork_id=1, name="rsv_weir"),
+        [
+            tabulated_rating_curve.Static(
+                level=[0.0, 1.5, 5],
+                flow_rate=[0.0, 45, 200],
+            )
+        ],
+    )
+
+    # Public water use
+    city = model.user_demand.add(
+        Node(9, Point(-0.75, -1), subnetwork_id=1, name="city"),
+        [
+            user_demand.Time(
+                # Total demand in m³/s
+                demand=[2.0, 2.3, 2.3, 2.4, 3, 3, 4, 3, 2.5, 2.2, 2.0, 2.0],
+                return_factor=0.4,
+                min_level=0,
+                priority=2,
+                time=pd.date_range(start="2022-01-01", periods=12, freq="MS"),
+            )
+        ],
+    )
+
+    # Industry
+    industry = model.user_demand.add(
+        Node(10, Point(0, -1.5), subnetwork_id=1, name="industry"),
+        [
+            user_demand.Time(
+                # Total demand in m³/s
+                demand=[4, 4, 4.5, 5, 5, 6, 7.5, 8, 5, 4, 3, 2.0],
+                return_factor=0.5,
+                min_level=0,
+                priority=1,
+                time=pd.date_range(start="2022-01-01", periods=12, freq="MS"),
+            )
+        ],
+    )
+
+    sea = model.terminal.add(Node(5, Point(-1.5, -3.0), subnetwork_id=1, name="sea"))
+
+    model.edge.add(main, reservoir, name="main")
+    model.edge.add(minor, conf, name="minor")
+    model.edge.add(reservoir, irr, name="irr supplied")
+    model.edge.add(irr, conf, name="irr drain")
+    model.edge.add(reservoir, city, name="city supplied")
+    model.edge.add(city, conf, name="city returnflow")
+    model.edge.add(reservoir, rsv_weir, name="rsv2weir")
+    model.edge.add(rsv_weir, conf, name="weir2conf")
+    model.edge.add(conf, tbr_conf, name="conf2tbr")
+    model.edge.add(level, reservoir)
+    model.edge.add(reservoir, industry, name="industry supplied")
+    model.edge.add(industry, conf, name="ind2conf")
+    model.edge.add(tbr_conf, sea, name="sea")
+
+    return model