From 118afd0ffd445bc5d638cdca331cd36f7ecc2f05 Mon Sep 17 00:00:00 2001
From: Huite Bootsma <huitebootsma@gmail.com>
Date: Thu, 25 Jan 2024 09:46:01 +0100
Subject: [PATCH 1/2] Add update_subgrid_level to BMI functions

---
 core/src/bmi.jl                                       |  5 +++++
 python/ribasim_api/tests/test_bmi.py                  | 11 +++++++++++
 python/ribasim_testmodels/ribasim_testmodels/basic.py | 11 ++++++++++-
 3 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/core/src/bmi.jl b/core/src/bmi.jl
index 7f11ee34d..83b31b2a2 100644
--- a/core/src/bmi.jl
+++ b/core/src/bmi.jl
@@ -633,6 +633,11 @@ function BMI.update_until(model::Model, time)::Model
     return model
 end
 
+function BMI.update_subgrid_level(model::Model)::Model
+    update_subgrid_level!(model.integrator)
+    return model
+end
+
 function BMI.get_value_ptr(model::Model, name::AbstractString)
     if name == "volume"
         model.integrator.u.storage
diff --git a/python/ribasim_api/tests/test_bmi.py b/python/ribasim_api/tests/test_bmi.py
index 4318eec81..01342a274 100644
--- a/python/ribasim_api/tests/test_bmi.py
+++ b/python/ribasim_api/tests/test_bmi.py
@@ -57,6 +57,17 @@ def test_update_until(libribasim, basic, tmp_path):
     assert actual_time == pytest.approx(expected_time)
 
 
+def test_update_subgrid_level(libribasim, basic, tmp_path):
+    basic.write(tmp_path / "ribasim.toml")
+    config_file = str(tmp_path / "ribasim.toml")
+    libribasim.initialize(config_file)
+    libribasim.update_subgrid_level()
+    level = libribasim.get_value_ptr("subgrid_level")
+    # The subgrid levels are initialized with NaN.
+    # After calling update, they should have regular values.
+    assert np.isfinite(level).all()
+
+
 def test_get_var_type(libribasim, basic, tmp_path):
     basic.write(tmp_path / "ribasim.toml")
     config_file = str(tmp_path / "ribasim.toml")
diff --git a/python/ribasim_testmodels/ribasim_testmodels/basic.py b/python/ribasim_testmodels/ribasim_testmodels/basic.py
index 657314b95..8d449d8c2 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/basic.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/basic.py
@@ -39,7 +39,16 @@ def basic_model() -> ribasim.Model:
     state = pd.DataFrame(
         data={"node_id": static["node_id"], "level": 0.04471158417652035}
     )
-    basin = ribasim.Basin(profile=profile, static=static, state=state)
+    # This is a 1:1 translation.
+    subgrid = pd.DataFrame(
+        data={
+            "node_id": profile["node_id"],
+            "subgrid_id": profile["node_id"],
+            "basin_level": profile["level"],
+            "subgrid_level": profile["level"],
+        }
+    )
+    basin = ribasim.Basin(profile=profile, static=static, state=state, subgrid=subgrid)
 
     # Setup linear resistance:
     linear_resistance = ribasim.LinearResistance(

From 8785aabb23df7c8514fe53163eeb740642bf3df2 Mon Sep 17 00:00:00 2001
From: Huite Bootsma <huitebootsma@gmail.com>
Date: Thu, 25 Jan 2024 10:38:32 +0100
Subject: [PATCH 2/2] Do not add to BMI namespace

---
 build/libribasim/src/libribasim.jl | 6 ++++++
 core/src/bmi.jl                    | 2 +-
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/build/libribasim/src/libribasim.jl b/build/libribasim/src/libribasim.jl
index ccaf3fb54..ee216288f 100644
--- a/build/libribasim/src/libribasim.jl
+++ b/build/libribasim/src/libribasim.jl
@@ -100,6 +100,12 @@ Base.@ccallable function update_until(time::Cdouble)::Cint
     end
 end
 
+Base.@ccallable function update_subgrid_level()::Cint
+    @try_c begin
+        Ribasim.update_subgrid_level(model)
+    end
+end
+
 Base.@ccallable function get_current_time(time::Ptr{Cdouble})::Cint
     @try_c begin
         t = BMI.get_current_time(model)
diff --git a/core/src/bmi.jl b/core/src/bmi.jl
index 83b31b2a2..9b9a7993d 100644
--- a/core/src/bmi.jl
+++ b/core/src/bmi.jl
@@ -633,7 +633,7 @@ function BMI.update_until(model::Model, time)::Model
     return model
 end
 
-function BMI.update_subgrid_level(model::Model)::Model
+function update_subgrid_level(model::Model)::Model
     update_subgrid_level!(model.integrator)
     return model
 end