From 76f4b32c651f518afe750a6960e1c418a8a1a09b Mon Sep 17 00:00:00 2001
From: Bart de Koning <74617371+SouthEndMusic@users.noreply.github.com>
Date: Wed, 17 Jan 2024 13:38:06 +0100
Subject: [PATCH] Fix looped subnetwork (#965)

Fixes https://github.com/Deltares/Ribasim/issues/692.

It now also runs fine with the cli (which I think I ran for the first
time ever ^^).
---
 .../ribasim_testmodels/__init__.py            |  6 ++----
 .../ribasim_testmodels/allocation.py          | 19 +++++++++++++------
 2 files changed, 15 insertions(+), 10 deletions(-)

diff --git a/python/ribasim_testmodels/ribasim_testmodels/__init__.py b/python/ribasim_testmodels/ribasim_testmodels/__init__.py
index a26ee232f..0ddc0cfb1 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/__init__.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/__init__.py
@@ -7,8 +7,8 @@
 import ribasim_testmodels
 from ribasim_testmodels.allocation import (
     allocation_example_model,
-    # looped_subnetwork_model,
     fractional_flow_subnetwork_model,
+    looped_subnetwork_model,
     minimal_subnetwork_model,
     subnetwork_model,
     user_model,
@@ -82,9 +82,7 @@
     "subnetwork_model",
     "minimal_subnetwork_model",
     "fractional_flow_subnetwork_model",
-    # Disable until this issue is resolved:
-    # https://github.com/Deltares/Ribasim/issues/692
-    # "looped_subnetwork_model",
+    "looped_subnetwork_model",
 ]
 
 # provide a mapping from model name to its constructor, so we can iterate over all models
diff --git a/python/ribasim_testmodels/ribasim_testmodels/allocation.py b/python/ribasim_testmodels/ribasim_testmodels/allocation.py
index c3416c4d4..4a218bd98 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/allocation.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/allocation.py
@@ -422,12 +422,15 @@ def looped_subnetwork_model():
         dtype=np.int64,
     )
     lines = node.geometry_from_connectivity(from_id, to_id)
+    allocation_network_id = len(from_id) * [None]
+    allocation_network_id[0] = 1
     edge = ribasim.Edge(
         df=gpd.GeoDataFrame(
             data={
                 "from_node_id": from_id,
                 "to_node_id": to_id,
                 "edge_type": len(from_id) * ["flow"],
+                "allocation_network_id": allocation_network_id,
             },
             geometry=lines,
             crs="EPSG:28992",
@@ -454,13 +457,13 @@ def looped_subnetwork_model():
         }
     )
 
-    state = pd.DataFrame(data={"node_id": [2, 7, 9, 15, 17, 21], "level": 1.0})
+    state = pd.DataFrame(data={"node_id": [2, 7, 9, 11, 15, 17, 21], "level": 1.0})
 
     basin = ribasim.Basin(profile=profile, static=static, state=state)
 
     # Setup the flow boundary:
     flow_boundary = ribasim.FlowBoundary(
-        static=pd.DataFrame(data={"node_id": [5], "flow_rate": [4.5]})
+        static=pd.DataFrame(data={"node_id": [5], "flow_rate": [4.5e-3]})
     )
 
     # Setup the users:
@@ -468,7 +471,7 @@ def looped_subnetwork_model():
         static=pd.DataFrame(
             data={
                 "node_id": [1, 12, 18, 20, 24],
-                "demand": 1.0,
+                "demand": 1.0e-3,
                 "return_factor": 0.9,
                 "min_level": 0.9,
                 "priority": [2, 1, 3, 3, 2],
@@ -481,8 +484,8 @@ def looped_subnetwork_model():
         static=pd.DataFrame(
             data={
                 "node_id": [6, 16],
-                "flow_rate": 4.0,
-                "max_flow_rate": 4.0,
+                "flow_rate": 4.0e-3,
+                "max_flow_rate": 4.0e-3,
             }
         )
     )
@@ -490,7 +493,7 @@ def looped_subnetwork_model():
     # Setup the outlets:
     outlet = ribasim.Outlet(
         static=pd.DataFrame(
-            data={"node_id": [3, 8, 10, 22], "flow_rate": 3.0, "max_flow_rate": 3.0}
+            data={"node_id": [3, 8, 10, 22], "flow_rate": 3.0e-3, "max_flow_rate": 3.0}
         )
     )
 
@@ -514,6 +517,9 @@ def looped_subnetwork_model():
         )
     )
 
+    # Setup allocation:
+    allocation = ribasim.Allocation(use_allocation=True, timestep=86400)
+
     model = ribasim.Model(
         network=ribasim.Network(node=node, edge=edge),
         basin=basin,
@@ -523,6 +529,7 @@ def looped_subnetwork_model():
         outlet=outlet,
         tabulated_rating_curve=rating_curve,
         terminal=terminal,
+        allocation=allocation,
         starttime="2020-01-01 00:00:00",
         endtime="2021-01-01 00:00:00",
     )