Catch residuals from OrdinaryDiffEq internals

core/src/Ribasim.jl

@@ -30,7 +30,7 @@ using DifferentiationInterface:
 using ForwardDiff: derivative as forward_diff
 
 # Algorithms for solving ODEs.
-using OrdinaryDiffEqCore: OrdinaryDiffEqCore, get_du
+using OrdinaryDiffEqCore: OrdinaryDiffEqCore, get_du, calculate_residuals!, ODE_DEFAULT_NORM
 import ForwardDiff
 
 # Interface for defining and solving the ODE problem of the physical layer.

core/src/callback.jl

@@ -122,7 +122,8 @@ Specifically, we first use all the inflows to update the mass of the Basins, rec
 the Basin concentration(s) and then remove the mass that is being lost to the outflows.
 """
 function update_cumulative_flows!(u, t, integrator)::Nothing
-    (; cache, p) = integrator
+    (; cache, p, opts) = integrator
+    (; internalnorm) = opts
     (; p_independent, p_mutable, time_dependent_cache) = p
     (; basin, flow_boundary, allocation, temp_convergence, convergence, ncalls) =
         p_independent
@@ -132,7 +133,9 @@ function update_cumulative_flows!(u, t, integrator)::Nothing
 
     # Update convergence measure
     if hasproperty(cache, :nlsolver)
-        @. temp_convergence = abs(cache.nlsolver.cache.atmp / u)
+        @. temp_convergence = abs(internalnorm.cache / (internalnorm.count * u))
+        internalnorm.cache .= 0
+        internalnorm.count = 0
         @inbounds for I in eachindex(temp_convergence)
             if !isfinite(temp_convergence[I])
                 temp_convergence[I] = zero(eltype(temp_convergence))
@@ -410,7 +413,6 @@ function save_flow(u, t, integrator)
         u_prev_saveat,
         convergence,
         ncalls,
-        node_id,
     ) = p.p_independent
     Δt = get_Δt(integrator)
     flow_mean = (u - u_prev_saveat) / Δt

core/src/model.jl

@@ -236,6 +236,7 @@ function Model(config::Config)::Model
     integrator = init(
         prob,
         alg;
+        internalnorm = NormWithCache(; cache = zero(u0)),
         progress = true,
         progress_name = "Simulating",
         progress_steps = 100,

core/src/util.jl

@@ -1255,3 +1255,43 @@ function should_skip_update_q(
 
     return control_type != control_type_
 end
+
+# Wrapper of ODE_DEFAULT_NORM used to dispatch on
+# for calculate_residuals!, also holding the residual
+# cache for output
+@kwdef mutable struct NormWithCache{T} <: Function
+    const cache::T
+    count::Int = 0
+end
+
+(::NormWithCache)(args...) = ODE_DEFAULT_NORM(args...)
+
+# Thin wrapper of `calculate_residuals!` to catch
+# the computed residuals
+function OrdinaryDiffEqCore.calculate_residuals!(
+    out,
+    ũ,
+    u₀,
+    u₁,
+    α,
+    ρ,
+    internalnorm::NormWithCache,
+    t,
+)
+    # Call the proper calculate_residuals!
+    output = invoke(
+        calculate_residuals!,
+        Tuple{Any, Any, Any, Any, Any, Any, Any, Any},
+        out,
+        ũ,
+        u₀,
+        u₁,
+        α,
+        ρ,
+        internalnorm,
+        t,
+    )
+    @. internalnorm.cache += abs(out)
+    internalnorm.count += 1
+    return output
+end