Add allocation subnetwork test model (#627)

Fixes #349.

core/src/solve.jl

@@ -411,7 +411,7 @@ return_factor: the factor in [0,1] of how much of the abstracted water is given
 min_level: The level of the source basin below which the user does not abstract
 priority: integer > 0, the lower the number the higher the priority of the users demand
 """
-struct User
+struct User <: AbstractParameterNode
     node_id::Vector{Int}
     active::BitVector
     demand::Vector{ScalarInterpolation}

core/src/utils.jl

@@ -737,7 +737,7 @@ Upstream basins always depend on themselves.
 function update_jac_prototype!(
     jac_prototype::SparseMatrixCSC{Float64, Int64},
     p::Parameters,
-    node::Union{Pump, Outlet, TabulatedRatingCurve},
+    node::Union{Pump, Outlet, TabulatedRatingCurve, User},
 )::Nothing
     (; basin, fractional_flow, connectivity) = p
     (; graph_flow) = connectivity

python/ribasim_testmodels/ribasim_testmodels/__init__.py

@@ -1,6 +1,6 @@
 __version__ = "0.2.0"
 
-from ribasim_testmodels.allocation import user_model
+from ribasim_testmodels.allocation import subnetwork_model, user_model
 from ribasim_testmodels.backwater import backwater_model
 from ribasim_testmodels.basic import (
     basic_model,
@@ -64,4 +64,5 @@
     "level_boundary_condition_model",
     "outlet_model",
     "user_model",
+    "subnetwork_model",
 ]

python/ribasim_testmodels/ribasim_testmodels/allocation.py

@@ -70,7 +70,7 @@ def user_model():
 
     basin = ribasim.Basin(profile=profile, static=static, state=state)
 
-    # Setup the user
+    # Setup the users:
     user = ribasim.User(
         static=pd.DataFrame(
             data={
@@ -122,3 +122,149 @@ def user_model():
     )
 
     return model
+
+
+def subnetwork_model():
+    """Create a user testmodel representing a subnetwork."""
+
+    # Setup the nodes:
+    xy = np.array(
+        [
+            (3.0, 1.0),  # 1: FlowBoundary
+            (2.0, 1.0),  # 2: Basin
+            (1.0, 1.0),  # 3: Outlet
+            (0.0, 1.0),  # 4: Terminal
+            (2.0, 2.0),  # 5: Pump
+            (2.0, 3.0),  # 6: Basin
+            (1.0, 3.0),  # 7: Outlet
+            (0.0, 3.0),  # 8: Basin
+            (2.0, 5.0),  # 9: Terminal
+            (2.0, 0.0),  # 10: User
+            (3.0, 3.0),  # 11: User
+            (0.0, 4.0),  # 12: User
+            (2.0, 4.0),  # 13: Outlet
+        ]
+    )
+    node_xy = gpd.points_from_xy(x=xy[:, 0], y=xy[:, 1])
+
+    node_type = [
+        "FlowBoundary",
+        "Basin",
+        "Outlet",
+        "Terminal",
+        "Pump",
+        "Basin",
+        "Outlet",
+        "Basin",
+        "Terminal",
+        "User",
+        "User",
+        "User",
+        "Outlet",
+    ]
+
+    # Make sure the feature id starts at 1: explicitly give an index.
+    node = ribasim.Node(
+        static=gpd.GeoDataFrame(
+            data={"type": node_type},
+            index=pd.Index(np.arange(len(xy)) + 1, name="fid"),
+            geometry=node_xy,
+            crs="EPSG:28992",
+        )
+    )
+
+    # Setup the edges:
+    from_id = np.array(
+        [1, 2, 3, 2, 2, 5, 6, 7, 6, 8, 6, 13, 10, 11, 12], dtype=np.int64
+    )
+    to_id = np.array([2, 3, 4, 10, 5, 6, 7, 8, 11, 12, 13, 9, 2, 6, 8], dtype=np.int64)
+    lines = ribasim.utils.geometry_from_connectivity(node, from_id, to_id)
+    edge = ribasim.Edge(
+        static=gpd.GeoDataFrame(
+            data={
+                "from_node_id": from_id,
+                "to_node_id": to_id,
+                "edge_type": len(from_id) * ["flow"],
+            },
+            geometry=lines,
+            crs="EPSG:28992",
+        )
+    )
+
+    # Setup the basins:
+    profile = pd.DataFrame(
+        data={"node_id": [2, 2, 6, 6, 8, 8], "area": 1000.0, "level": 3 * [0.0, 1.0]}
+    )
+
+    static = pd.DataFrame(
+        data={
+            "node_id": [2, 6, 8],
+            "drainage": 0.0,
+            "potential_evaporation": 0.0,
+            "infiltration": 0.0,
+            "precipitation": 0.0,
+            "urban_runoff": 0.0,
+        }
+    )
+
+    # Setup the flow boundary:
+    flow_boundary = ribasim.FlowBoundary(
+        static=pd.DataFrame(data={"node_id": [1], "flow_rate": [1e-3]})
+    )
+
+    state = pd.DataFrame(data={"node_id": [2, 6, 8], "level": 1.0})
+
+    basin = ribasim.Basin(profile=profile, static=static, state=state)
+
+    # Setup the users:
+    user = ribasim.User(
+        static=pd.DataFrame(
+            data={
+                "node_id": [10, 11, 12],
+                "demand": [1.0, 2.0, 4.0],
+                "return_factor": 0.9,
+                "min_level": 0.9,
+                "priority": [2, 2, 1],
+            }
+        )
+    )
+
+    # Setup the pump:
+    pump = ribasim.Pump(
+        static=pd.DataFrame(
+            data={
+                "node_id": [5],
+                "flow_rate": [4.0],
+            }
+        )
+    )
+
+    # Setup the outlets:
+    outlet = ribasim.Outlet(
+        static=pd.DataFrame(data={"node_id": [3, 7, 13], "flow_rate": 3.0})
+    )
+
+    # Setup the terminal:
+    terminal = ribasim.Terminal(
+        static=pd.DataFrame(
+            data={
+                "node_id": [4, 9],
+            }
+        )
+    )
+
+    model = ribasim.Model(
+        modelname="subnetwork",
+        node=node,
+        edge=edge,
+        basin=basin,
+        user=user,
+        flow_boundary=flow_boundary,
+        pump=pump,
+        outlet=outlet,
+        terminal=terminal,
+        starttime="2020-01-01 00:00:00",
+        endtime="2021-01-01 00:00:00",
+    )
+
+    return model