From 47b4fa7eccd307408dafb862545c0bce31f36299 Mon Sep 17 00:00:00 2001
From: Maarten Pronk <git@evetion.nl>
Date: Thu, 14 Nov 2024 16:58:44 +0100
Subject: [PATCH] Add some basic concentration tests. (#1934)

Ideally we merge this before #1933 and #1929, so we know they don't
break anything.

---------

Co-authored-by: Martijn Visser <mgvisser@gmail.com>
---
 core/src/util.jl                                             | 1 +
 core/test/run_models_test.jl                                 | 5 +++++
 python/ribasim_testmodels/ribasim_testmodels/basic.py        | 3 ++-
 .../ribasim_testmodels/discrete_control.py                   | 3 ++-
 4 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/core/src/util.jl b/core/src/util.jl
index d12498c9b..0ccfddb28 100644
--- a/core/src/util.jl
+++ b/core/src/util.jl
@@ -359,6 +359,7 @@ function FlatVector(saveval::Vector{<:SavedFlow}, sym::Symbol)
     v = isempty(saveval) ? Vector{Float64}[] : getfield.(saveval, sym)
     FlatVector(v)
 end
+FlatVector(v::Vector{Matrix{Float64}}) = FlatVector(vec.(v))
 
 """
 Function that goes smoothly from 0 to 1 in the interval [0,threshold],
diff --git a/core/test/run_models_test.jl b/core/test/run_models_test.jl
index fc6c5e66d..d7d1bf71b 100644
--- a/core/test/run_models_test.jl
+++ b/core/test/run_models_test.jl
@@ -222,6 +222,11 @@ end
 
     # flows are recorded at the end of each period, and are undefined at the start
     @test unique(table.time) == Ribasim.datetimes(model)[1:(end - 1)]
+
+    @test isfile(joinpath(dirname(toml_path), "results/concentration.arrow"))
+    table = Ribasim.concentration_table(model)
+    @test "Continuity" in table.substance
+    @test all(isapprox.(table.concentration[table.substance .== "Continuity"], 1.0))
 end
 
 @testitem "basic arrow model" begin
diff --git a/python/ribasim_testmodels/ribasim_testmodels/basic.py b/python/ribasim_testmodels/ribasim_testmodels/basic.py
index f431150a7..8618631de 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/basic.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/basic.py
@@ -5,7 +5,7 @@
 import numpy as np
 import pandas as pd
 import ribasim
-from ribasim.config import Node
+from ribasim.config import Experimental, Node
 from ribasim.input_base import TableModel
 from ribasim.nodes import (
     basin,
@@ -26,6 +26,7 @@ def basic_model() -> ribasim.Model:
         starttime="2020-01-01",
         endtime="2021-01-01",
         crs="EPSG:28992",
+        experimental=Experimental(concentration=True),
     )
     model.logging = ribasim.Logging(verbosity="debug")
 
diff --git a/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py b/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py
index 745efc563..c325fbfb5 100644
--- a/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py
+++ b/python/ribasim_testmodels/ribasim_testmodels/discrete_control.py
@@ -1,6 +1,6 @@
 import numpy as np
 import pandas as pd
-from ribasim.config import Node
+from ribasim.config import Experimental, Node
 from ribasim.model import Model, Solver
 from ribasim.nodes import (
     basin,
@@ -604,6 +604,7 @@ def continuous_concentration_condition_model() -> Model:
         endtime="2020-02-01",
         crs="EPSG:28992",
         solver=Solver(saveat=86400 / 8),
+        experimental=Experimental(concentration=True),
     )
 
     basi = model.basin.add(