From 0ba441b58f1108c3f2ea0c6f2547211af8ff7e81 Mon Sep 17 00:00:00 2001
From: Bart de Koning <74617371+SouthEndMusic@users.noreply.github.com>
Date: Tue, 21 May 2024 17:22:56 +0200
Subject: [PATCH] Allow `missing` as static user demand (#1493)

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

---------

Co-authored-by: Martijn Visser <mgvisser@gmail.com>
---
 core/src/read.jl                                          | 5 +++--
 core/src/schema.jl                                        | 2 +-
 core/test/allocation_test.jl                              | 7 ++++---
 python/ribasim/ribasim/schemas.py                         | 2 +-
 .../ribasim_testmodels/ribasim_testmodels/allocation.py   | 8 ++++++++
 5 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/core/src/read.jl b/core/src/read.jl
index 9c2ac0517..554c3c0db 100644
--- a/core/src/read.jl
+++ b/core/src/read.jl
@@ -794,8 +794,9 @@ function user_demand_static!(
 
         for row in group
             priority_idx = findsorted(priorities, row.priority)
-            demand_itp[user_demand_idx][priority_idx].u .= row.demand
-            demand[user_demand_idx, priority_idx] = row.demand
+            demand_row = coalesce(row.demand, 0.0)
+            demand_itp[user_demand_idx][priority_idx].u .= demand_row
+            demand[user_demand_idx, priority_idx] = demand_row
         end
     end
     return nothing
diff --git a/core/src/schema.jl b/core/src/schema.jl
index 1b6c8b42d..e36d4d128 100644
--- a/core/src/schema.jl
+++ b/core/src/schema.jl
@@ -273,7 +273,7 @@ end
 @version UserDemandStaticV1 begin
     node_id::Int32
     active::Union{Missing, Bool}
-    demand::Float64
+    demand::Union{Missing, Float64}
     return_factor::Float64
     min_level::Float64
     priority::Int32
diff --git a/core/test/allocation_test.jl b/core/test/allocation_test.jl
index 1171743b7..ca8a4285e 100644
--- a/core/test/allocation_test.jl
+++ b/core/test/allocation_test.jl
@@ -177,9 +177,10 @@ end
     # Subnetworks interpreted as user_demands require variables and constraints to
     # support absolute value expressions in the objective function
     problem = allocation_model_main_network.problem
-    @test problem[:F_abs_user_demand].axes[1] == NodeID.(:Pump, [11, 24, 38])
-    @test problem[:abs_positive_user_demand].axes[1] == NodeID.(:Pump, [11, 24, 38])
-    @test problem[:abs_negative_user_demand].axes[1] == NodeID.(:Pump, [11, 24, 38])
+    node_ids = [NodeID(:UserDemand, 60), NodeID.(:Pump, [11, 24, 38])...]
+    @test problem[:F_abs_user_demand].axes[1] == node_ids
+    @test problem[:abs_positive_user_demand].axes[1] == node_ids
+    @test problem[:abs_negative_user_demand].axes[1] == node_ids
 
     # In each subnetwork, the connection from the main network to the subnetwork is
     # interpreted as a source
diff --git a/python/ribasim/ribasim/schemas.py b/python/ribasim/ribasim/schemas.py
index 3686b2398..cf1bb4477 100644
--- a/python/ribasim/ribasim/schemas.py
+++ b/python/ribasim/ribasim/schemas.py
@@ -246,7 +246,7 @@ class TerminalStaticSchema(_BaseSchema):
 class UserDemandStaticSchema(_BaseSchema):
     node_id: Series[Int32] = pa.Field(nullable=False, default=0)
     active: Series[pa.BOOL] = pa.Field(nullable=True)
-    demand: Series[float] = pa.Field(nullable=False)
+    demand: Series[float] = pa.Field(nullable=True)
     return_factor: Series[float] = pa.Field(nullable=False)
     min_level: Series[float] = pa.Field(nullable=False)
     priority: Series[Int32] = pa.Field(nullable=False, default=0)
diff --git a/python/ribasim_testmodels/ribasim_testmodels/allocation.py b/python/ribasim_testmodels/ribasim_testmodels/allocation.py
index 746f8dd4b..8ce6665c5 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/allocation.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/allocation.py
@@ -790,6 +790,12 @@ def main_network_with_subnetworks_model() -> Model:
         ],
     )
 
+    # Missing demand
+    model.user_demand.add(
+        Node(60, Point(21, -1), subnetwork_id=1),
+        [user_demand.Static(return_factor=[0.9], priority=2, min_level=0.0)],
+    )
+
     model.edge.add(model.flow_boundary[1], model.basin[2], subnetwork_id=1)
     model.edge.add(model.basin[2], model.linear_resistance[3])
     model.edge.add(model.linear_resistance[3], model.basin[4])
@@ -858,6 +864,8 @@ def main_network_with_subnetworks_model() -> Model:
     model.edge.add(model.basin[2], model.pump[11], subnetwork_id=3)
     model.edge.add(model.basin[6], model.pump[24], subnetwork_id=5)
     model.edge.add(model.basin[10], model.pump[38], subnetwork_id=7)
+    model.edge.add(model.basin[8], model.user_demand[60])
+    model.edge.add(model.user_demand[60], model.basin[8])
 
     return model