Pass the integrator and the stage buffer to stage limiters (#4356)

Author: singhharsh1708

lib/OrdinaryDiffEqLowOrderRK/Project.toml

@@ -1,7 +1,7 @@
 name = "OrdinaryDiffEqLowOrderRK"
 uuid = "1344f307-1e59-4825-a18e-ace9aa3fa4c6"
 authors = ["ParamThakkar123 <paramthakkar864@gmail.com>"]
-version = "2.2.3"
+version = "2.2.4"
 
 [deps]
 MuladdMacro = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221"

lib/OrdinaryDiffEqLowOrderRK/src/fixed_timestep_perform_step.jl

@@ -219,7 +219,7 @@ end
     stage_limiter! = integrator.opts.stage_limiter!
     halfdt = dt / 2
     @.. broadcast = false thread = thread tmp = uprev + halfdt * fsalfirst
-    stage_limiter!(k, tmp, p, t + halfdt)
+    stage_limiter!(tmp, integrator, p, t + halfdt)
     f(k, tmp, p, t + halfdt)
     OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
     @.. broadcast = false thread = thread u = uprev + dt * k

lib/OrdinaryDiffEqQPRK/Project.toml

@@ -1,7 +1,7 @@
 name = "OrdinaryDiffEqQPRK"
 uuid = "04162be5-8125-4266-98ed-640baecc6514"
 authors = ["ParamThakkar123 <paramthakkar864@gmail.com>"]
-version = "2.2.0"
+version = "2.2.1"
 
 [deps]
 MuladdMacro = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221"

lib/OrdinaryDiffEqQPRK/src/qprk_perform_step.jl

@@ -134,7 +134,7 @@ end
     stage_limiter!(tmp, integrator, p, t + d4 * dt)
     f(k4, tmp, p, t + d4 * dt)
     @.. broadcast = false thread = thread tmp = uprev + dt * (b51 * k1 + b53 * k3 + b54 * k4)
-    stage_limiter!(uprev, integrator, p, t + d5 * dt)
+    stage_limiter!(tmp, integrator, p, t + d5 * dt)
     f(k5, tmp, p, t + d5 * dt)
     @.. broadcast = false thread = thread tmp = uprev + dt * (b61 * k1 + b64 * k4 + b65 * k5)
     stage_limiter!(tmp, integrator, p, t + d6 * dt)
@@ -212,7 +212,7 @@ end
             + b16_10 * k10 + b16_11 * k11 + b16_12 * k12
             + b16_13 * k13 + b16_14 * k14
     )
-    stage_limiter!(u, integrator, p, t + dt)
+    stage_limiter!(tmp, integrator, p, t + dt)
     f(k16, tmp, p, t + dt)
 
     OrdinaryDiffEqCore.increment_nf!(integrator.stats, 16)

lib/OrdinaryDiffEqTsit5/Project.toml

@@ -1,6 +1,6 @@
 name = "OrdinaryDiffEqTsit5"
 uuid = "b1df2697-797e-41e3-8120-5422d3b24e4a"
-version = "2.1.3"
+version = "2.1.4"
 authors = ["ParamThakkar123 <paramthakkar864@gmail.com>"]
 
 [deps]

lib/OrdinaryDiffEqTsit5/src/tsit_perform_step.jl

@@ -210,24 +210,24 @@ end
     stage_limiter! = integrator.opts.stage_limiter!
     a = dt * a21
     @.. broadcast = false thread = thread tmp = uprev + a * k1
-    stage_limiter!(tmp, f, p, t + c1 * dt)
+    stage_limiter!(tmp, integrator, p, t + c1 * dt)
     f(k2, tmp, p, t + c1 * dt)
     @.. broadcast = false thread = thread tmp = uprev + dt * (a31 * k1 + a32 * k2)
-    stage_limiter!(tmp, f, p, t + c2 * dt)
+    stage_limiter!(tmp, integrator, p, t + c2 * dt)
     f(k3, tmp, p, t + c2 * dt)
     @.. broadcast = false thread = thread tmp = uprev + dt * (a41 * k1 + a42 * k2 + a43 * k3)
-    stage_limiter!(tmp, f, p, t + c3 * dt)
+    stage_limiter!(tmp, integrator, p, t + c3 * dt)
     f(k4, tmp, p, t + c3 * dt)
     @.. broadcast = false thread = thread tmp = uprev +
         dt * (a51 * k1 + a52 * k2 + a53 * k3 + a54 * k4)
-    stage_limiter!(tmp, f, p, t + c4 * dt)
+    stage_limiter!(tmp, integrator, p, t + c4 * dt)
     f(k5, tmp, p, t + c4 * dt)
     @.. broadcast = false thread = thread tmp = uprev +
         dt * (
         a61 * k1 + a62 * k2 + a63 * k3 + a64 * k4 +
             a65 * k5
     )
-    stage_limiter!(tmp, f, p, t + dt)
+    stage_limiter!(tmp, integrator, p, t + dt)
     f(k6, tmp, p, t + dt)
     @.. broadcast = false thread = thread u = uprev +
         dt * (

test/Integrators_I/step_limiter_test.jl

@@ -191,3 +191,63 @@ end
         @test_throws ErrorException solve(prob, alg, dt = 0.1; stage_limiter = slim!)
     end
 end
+
+# The stage limiter contract is `limiter!(u, integrator, p, t)`
+# (`OrdinaryDiffEqCore.trivial_limiter!`). The limiters above ignore their
+# arguments, so they accept whatever is handed to them; that is how `Tsit5` shipped
+# with the `ODEFunction` in the integrator slot, `Midpoint` with a stage buffer, and
+# `QPRK98` limiting `uprev`. These testsets pin down which objects get passed rather
+# than how many times the limiter is called.
+const LIMITER_ALGS = [
+    Euler, Heun, Ralston, Midpoint, RK4, BS3, OwrenZen3, DP5, Tsit5,
+    Vern6, Vern9, DP8, TanYam7, TsitPap8, QPRK98,
+    SSPRK22, SSPRK43, SSPRK104, SSPRK932,
+    CarpenterKennedy2N54, ORK256, RDPK3Sp35, NDBLSRK124,
+    Rosenbrock23, ROS3P, Rodas4, Rodas5P,
+]
+
+@testset "stage limiter receives the integrator" begin
+    prob = ODEProblem((du, u, p, t) -> du .= u, [1.0, 1.0], (0.0, 1.0))
+    for A in LIMITER_ALGS
+        ok = Ref(true)
+        calls = Ref(0)
+        limiter! = function (u, integrator, p, t)
+            calls[] += 1
+            integrator isa SciMLBase.DEIntegrator || (ok[] = false)
+            return nothing
+        end
+        solve(prob, A(), dt = 0.1; stage_limiter = limiter!)
+        @test calls[] > 0
+        @test ok[]
+    end
+end
+
+# The limiter mutates whatever it is given, so handing it `uprev` corrupts every
+# later stage and the step update. `QPRK98` stage 5 did exactly that.
+@testset "stage limiter is never handed uprev" begin
+    prob = ODEProblem((du, u, p, t) -> du .= u, [1.0, 1.0], (0.0, 1.0))
+    for A in LIMITER_ALGS
+        ok = Ref(true)
+        limiter! = function (u, integrator, p, t)
+            u === integrator.uprev && (ok[] = false)
+            return nothing
+        end
+        solve(prob, A(), dt = 0.1; stage_limiter = limiter!)
+        @test ok[]
+    end
+end
+
+# End-to-end: with `QPRK98` stage 5 clamping `uprev` in place, `uprev` changed
+# value partway through a single step.
+@testset "uprev is stable across one step" begin
+    prob = ODEProblem((du, u, p, t) -> du .= u, [-1.0], (0.0, 0.1))
+    seen = Float64[]
+    positivity! = function (u, integrator, p, t)
+        push!(seen, integrator.uprev[1])
+        @. u = max(u, 0.0)
+        return nothing
+    end
+    solve(prob, QPRK98(), dt = 0.1, adaptive = false; stage_limiter = positivity!)
+    @test !isempty(seen)
+    @test all(==(seen[1]), seen)
+end