Attempt at #1084

Deltares · Mar 26, 2024 · 6ab85b3 · 6ab85b3
1 parent bf0dfe0
commit 6ab85b3
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 3 deletions.
diff --git a/core/src/bmi.jl b/core/src/bmi.jl
@@ -36,10 +36,14 @@ function BMI.get_value_ptr(model::Model, name::AbstractString)
         model.integrator.u.storage
     elseif name == "basin.level"
         get_tmp(model.integrator.p.basin.current_level, 0)
-    elseif name == "basin.infiltration"
+    elseif name == "basin.infiltration_instantaneous"
         get_tmp(model.integrator.p.basin.vertical_flux, 0).infiltration
-    elseif name == "basin.drainage"
+    elseif name == "basin.drainage_instantaneous"
         get_tmp(model.integrator.p.basin.vertical_flux, 0).drainage
+    elseif name == "basin.infiltration_integrated"
+        model.integrator.p.basin.vertical_flux_bmi.infiltration
+    elseif name == "basin.drainage_integrated"
+        model.integrator.p.basin.vertical_flux_bmi.drainage
     elseif name == "basin.subgrid_level"
         model.integrator.p.subgrid.level
     elseif name == "user_demand.demand"

diff --git a/core/src/callback.jl b/core/src/callback.jl
@@ -111,7 +111,8 @@ function integrate_flows!(u, t, integrator)::Nothing
     (; p, dt) = integrator
     (; graph, user_demand, basin) = p
     (; flow, flow_dict, flow_prev, flow_integrated) = graph[]
-    (; vertical_flux, vertical_flux_prev, vertical_flux_integrated) = basin
+    (; vertical_flux, vertical_flux_prev, vertical_flux_integrated, vertical_flux_bmi) =
+        basin
     flow = get_tmp(flow, 0)
     vertical_flux = get_tmp(vertical_flux, 0)
     if !isempty(flow_prev) && isnan(flow_prev[1])
@@ -121,6 +122,7 @@ function integrate_flows!(u, t, integrator)::Nothing
 
     @. flow_integrated += 0.5 * (flow + flow_prev) * dt
     @. vertical_flux_integrated += 0.5 * (vertical_flux + vertical_flux_prev) * dt
+    @. vertical_flux_bmi += 0.5 * (vertical_flux + vertical_flux_prev) * dt
 
     for (i, id) in enumerate(user_demand.node_id)
         src_id = inflow_id(graph, id)

diff --git a/core/src/parameter.jl b/core/src/parameter.jl
@@ -158,6 +158,7 @@ struct Basin{T, C, V1, V2, V3} <: AbstractParameterNode
     vertical_flux::V2
     vertical_flux_prev::V3
     vertical_flux_integrated::V3
+    vertical_flux_bmi::V3
     # Cache this to avoid recomputation
     current_level::T
     current_area::T
@@ -176,6 +177,7 @@ struct Basin{T, C, V1, V2, V3} <: AbstractParameterNode
         vertical_flux::V2,
         vertical_flux_prev::V3,
         vertical_flux_integrated,
+        vertical_flux_bmi,
         current_level::T,
         current_area::T,
         area,
@@ -192,6 +194,7 @@ struct Basin{T, C, V1, V2, V3} <: AbstractParameterNode
             vertical_flux,
             vertical_flux_prev,
             vertical_flux_integrated,
+            vertical_flux_bmi,
             current_level,
             current_area,
             area,

diff --git a/core/src/read.jl b/core/src/read.jl
@@ -515,6 +515,7 @@ function Basin(db::DB, config::Config, chunk_sizes::Vector{Int})::Basin
     )
     vertical_flux_prev = zero(vertical_flux)
     vertical_flux_integrated = zero(vertical_flux)
+    vertical_flux_bmi = zero(vertical_flux)
 
     if config.solver.autodiff
         current_level = DiffCache(current_level, chunk_sizes)
@@ -530,6 +531,7 @@ function Basin(db::DB, config::Config, chunk_sizes::Vector{Int})::Basin
         vertical_flux,
         vertical_flux_prev,
         vertical_flux_integrated,
+        vertical_flux_bmi,
         current_level,
         current_area,
         area,