From 2be7d93fd90e6b7880910f7bca65b88301661b11 Mon Sep 17 00:00:00 2001
From: Bart de Koning <74617371+SouthEndMusic@users.noreply.github.com>
Date: Tue, 23 Apr 2024 22:35:21 +0200
Subject: [PATCH] Negative storage check (#1425)

Fixes https://github.com/Deltares/Ribasim/issues/1406.
---
 core/src/callback.jl         | 21 +++++++++++++++++++++
 core/test/allocation_test.jl |  2 +-
 core/test/validation_test.jl | 20 ++++++++++++++++++++
 3 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/core/src/callback.jl b/core/src/callback.jl
index 4344e2af6..51dd3d029 100644
--- a/core/src/callback.jl
+++ b/core/src/callback.jl
@@ -13,6 +13,9 @@ function create_callbacks(
     (; starttime, basin, tabulated_rating_curve, discrete_control) = parameters
     callbacks = SciMLBase.DECallback[]
 
+    negative_storage_cb = FunctionCallingCallback(check_negative_storage)
+    push!(callbacks, negative_storage_cb)
+
     integrating_flows_cb = FunctionCallingCallback(integrate_flows!; func_start = false)
     push!(callbacks, integrating_flows_cb)
 
@@ -65,6 +68,24 @@ function create_callbacks(
     return callback, saved
 end
 
+function check_negative_storage(u, t, integrator)::Nothing
+    (; basin) = integrator.p
+    (; node_id) = basin
+    errors = false
+    for (i, id) in enumerate(node_id)
+        if u.storage[i] < 0
+            @error "Negative storage detected in $id"
+            errors = true
+        end
+    end
+
+    if errors
+        t_datetime = datetime_since(integrator.t, integrator.p.starttime)
+        error("Negative storages found at $t_datetime.")
+    end
+    return nothing
+end
+
 """
 Integrate flows over the last timestep
 """
diff --git a/core/test/allocation_test.jl b/core/test/allocation_test.jl
index f5ecfa916..023f9dee4 100644
--- a/core/test/allocation_test.jl
+++ b/core/test/allocation_test.jl
@@ -68,7 +68,7 @@ end
         @__DIR__,
         "../../generated_testmodels/fractional_flow_subnetwork/ribasim.toml",
     )
-    model = Ribasim.BMI.initialize(Ribasim.Model, toml_path)
+    model = Ribasim.Model(toml_path)
     problem = model.integrator.p.allocation.allocation_models[1].problem
     F = problem[:F]
     constraints_fractional_flow = problem[:fractional_flow]
diff --git a/core/test/validation_test.jl b/core/test/validation_test.jl
index 3a78691c7..8ca4ac547 100644
--- a/core/test/validation_test.jl
+++ b/core/test/validation_test.jl
@@ -438,3 +438,23 @@ end
     @test logger.logs[1].message ==
           "Demand of UserDemand #1 with priority 1 should be non-negative"
 end
+
+@testitem "negative storage" begin
+    import BasicModelInterface as BMI
+    toml_path =
+        normpath(@__DIR__, "../../generated_testmodels/linear_resistance/ribasim.toml")
+    @test ispath(toml_path)
+    dt = 1e10
+
+    config = Ribasim.Config(
+        toml_path;
+        solver_algorithm = "Euler",
+        solver_dt = dt,
+        solver_saveat = Inf,
+    )
+    model = Ribasim.Model(config)
+    @test_throws "Negative storages found at 2021-01-01T00:00:00." BMI.update_until(
+        model,
+        dt,
+    )
+end