LowStorageRK: unify 3Sp + 3SpFSAL families on generic perform_step

Author: Your Name

lib/OrdinaryDiffEqLowStorageRK/src/generic_low_storage_perform_step.jl

@@ -137,3 +137,158 @@ end
     OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
     return nothing
 end
+
+@muladd function _perform_step_oop!(integrator, tab::LowStorageRK3SpConstantCache)
+    (; t, dt, uprev, u, f, p) = integrator
+    (; γ12end, γ22end, γ32end, δ2end, β1, β2end, c2end, bhat1, bhat2end) = tab
+
+    integrator.fsalfirst = f(uprev, p, t)
+    OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
+    integrator.k[1] = integrator.fsalfirst
+    tmp = uprev
+    u = tmp + β1 * dt * integrator.fsalfirst
+    utilde = u
+    if integrator.opts.adaptive
+        utilde = bhat1 * dt * integrator.fsalfirst
+    end
+
+    for i in eachindex(γ12end)
+        k = f(u, p, t + c2end[i] * dt)
+        OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
+        tmp = tmp + δ2end[i] * u
+        u = γ12end[i] * u + γ22end[i] * tmp + γ32end[i] * uprev + β2end[i] * dt * k
+        if integrator.opts.adaptive
+            utilde = utilde + bhat2end[i] * dt * k
+        end
+    end
+
+    if integrator.opts.adaptive
+        atmp = calculate_residuals(
+            utilde, uprev, u, integrator.opts.abstol,
+            integrator.opts.reltol, integrator.opts.internalnorm, t
+        )
+        OrdinaryDiffEqCore.set_EEst!(integrator, integrator.opts.internalnorm(atmp, t))
+    end
+
+    integrator.u = u
+    return nothing
+end
+
+@muladd function _perform_step_iip!(integrator, cache, tab::LowStorageRK3SpConstantCache)
+    (; t, dt, uprev, u, f, p) = integrator
+    (; k, tmp, utilde, atmp, stage_limiter!, step_limiter!, thread) = cache
+    (; γ12end, γ22end, γ32end, δ2end, β1, β2end, c2end, bhat1, bhat2end) = tab
+
+    f(integrator.fsalfirst, uprev, p, t)
+    OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
+    @.. broadcast = false thread = thread tmp = uprev
+    @.. broadcast = false thread = thread u = tmp + β1 * dt * integrator.fsalfirst
+    if integrator.opts.adaptive
+        @.. broadcast = false thread = thread utilde = bhat1 * dt * integrator.fsalfirst
+    end
+
+    for i in eachindex(γ12end)
+        stage_limiter!(u, integrator, p, t + c2end[i] * dt)
+        f(k, u, p, t + c2end[i] * dt)
+        OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
+        @.. broadcast = false thread = thread tmp = tmp + δ2end[i] * u
+        @.. broadcast = false thread = thread u = γ12end[i] * u + γ22end[i] * tmp +
+            γ32end[i] * uprev + β2end[i] * dt * k
+        if integrator.opts.adaptive
+            @.. broadcast = false thread = thread utilde = utilde + bhat2end[i] * dt * k
+        end
+    end
+
+    stage_limiter!(u, integrator, p, t + dt)
+    step_limiter!(u, integrator, p, t + dt)
+
+    if integrator.opts.adaptive
+        calculate_residuals!(
+            atmp, utilde, uprev, u, integrator.opts.abstol,
+            integrator.opts.reltol, integrator.opts.internalnorm, t,
+            thread
+        )
+        OrdinaryDiffEqCore.set_EEst!(integrator, integrator.opts.internalnorm(atmp, t))
+    end
+    return nothing
+end
+
+@muladd function _perform_step_oop!(integrator, tab::LowStorageRK3SpFSALConstantCache)
+    (; t, dt, uprev, u, f, p) = integrator
+    (; γ12end, γ22end, γ32end, δ2end, β1, β2end, c2end, bhat1, bhat2end, bhatfsal) = tab
+
+    tmp = uprev
+    u = tmp + β1 * dt * integrator.fsalfirst
+    utilde = u
+    if integrator.opts.adaptive
+        utilde = bhat1 * dt * integrator.fsalfirst
+    end
+
+    for i in eachindex(γ12end)
+        k = f(u, p, t + c2end[i] * dt)
+        OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
+        tmp = tmp + δ2end[i] * u
+        u = γ12end[i] * u + γ22end[i] * tmp + γ32end[i] * uprev + β2end[i] * dt * k
+        if integrator.opts.adaptive
+            utilde = utilde + bhat2end[i] * dt * k
+        end
+    end
+
+    integrator.fsallast = f(u, p, t + dt)
+    OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
+
+    if integrator.opts.adaptive
+        utilde = utilde + bhatfsal * dt * integrator.fsallast
+        atmp = calculate_residuals(
+            utilde, uprev, u, integrator.opts.abstol,
+            integrator.opts.reltol, integrator.opts.internalnorm, t
+        )
+        OrdinaryDiffEqCore.set_EEst!(integrator, integrator.opts.internalnorm(atmp, t))
+    end
+
+    integrator.k[1] = integrator.fsalfirst
+    integrator.k[2] = integrator.fsallast
+    integrator.u = u
+    return nothing
+end
+
+@muladd function _perform_step_iip!(integrator, cache, tab::LowStorageRK3SpFSALConstantCache)
+    (; t, dt, uprev, u, f, p) = integrator
+    (; k, tmp, utilde, atmp, stage_limiter!, step_limiter!, thread) = cache
+    (; γ12end, γ22end, γ32end, δ2end, β1, β2end, c2end, bhat1, bhat2end, bhatfsal) = tab
+
+    @.. broadcast = false thread = thread tmp = uprev
+    @.. broadcast = false thread = thread u = tmp + β1 * dt * integrator.fsalfirst
+    if integrator.opts.adaptive
+        @.. broadcast = false thread = thread utilde = bhat1 * dt * integrator.fsalfirst
+    end
+
+    for i in eachindex(γ12end)
+        stage_limiter!(u, integrator, p, t + c2end[i] * dt)
+        f(k, u, p, t + c2end[i] * dt)
+        OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
+        @.. broadcast = false thread = thread tmp = tmp + δ2end[i] * u
+        @.. broadcast = false thread = thread u = γ12end[i] * u + γ22end[i] * tmp +
+            γ32end[i] * uprev + β2end[i] * dt * k
+        if integrator.opts.adaptive
+            @.. broadcast = false thread = thread utilde = utilde + bhat2end[i] * dt * k
+        end
+    end
+
+    stage_limiter!(u, integrator, p, t + dt)
+    step_limiter!(u, integrator, p, t + dt)
+
+    f(k, u, p, t + dt)
+    OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
+
+    if integrator.opts.adaptive
+        @.. broadcast = false thread = thread utilde = utilde + bhatfsal * dt * k
+        calculate_residuals!(
+            atmp, utilde, uprev, u, integrator.opts.abstol,
+            integrator.opts.reltol, integrator.opts.internalnorm, t,
+            thread
+        )
+        OrdinaryDiffEqCore.set_EEst!(integrator, integrator.opts.internalnorm(atmp, t))
+    end
+    return nothing
+end

lib/OrdinaryDiffEqLowStorageRK/src/low_storage_rk_perform_step.jl

@@ -77,223 +77,56 @@ function perform_step!(integrator, cache::LowStorageRK3SCache, repeat_step = fal
     return _perform_step_iip!(integrator, cache, cache.tab)
 end
 
-# 3S+ low storage methods: 3S methods adding another memory location for the embedded method (non-FSAL version)
 function initialize!(integrator, cache::LowStorageRK3SpConstantCache)
-    integrator.fsalfirst = integrator.f(integrator.uprev, integrator.p, integrator.t) # Pre-start fsal
+    integrator.fsalfirst = integrator.f(integrator.uprev, integrator.p, integrator.t)
     OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
     integrator.kshortsize = 1
     integrator.k = typeof(integrator.k)(undef, integrator.kshortsize)
-
-    # Avoid undefined entries if k is an array of arrays
     integrator.fsallast = zero(integrator.fsalfirst)
     return integrator.k[1] = integrator.fsalfirst
 end
 
-@muladd function perform_step!(
-        integrator, cache::LowStorageRK3SpConstantCache,
-        repeat_step = false
-    )
-    (; t, dt, uprev, u, f, p) = integrator
-    (; γ12end, γ22end, γ32end, δ2end, β1, β2end, c2end, bhat1, bhat2end) = cache
-
-    # u1
-    integrator.fsalfirst = f(uprev, p, t)
-    OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
-    integrator.k[1] = integrator.fsalfirst
-    tmp = uprev
-    u = tmp + β1 * dt * integrator.fsalfirst
-    # Initialize utilde for JET
-    utilde = u
-    if integrator.opts.adaptive
-        utilde = bhat1 * dt * integrator.fsalfirst
-    end
-
-    # other stages
-    for i in eachindex(γ12end)
-        k = f(u, p, t + c2end[i] * dt)
-        OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
-        tmp = tmp + δ2end[i] * u
-        u = γ12end[i] * u + γ22end[i] * tmp + γ32end[i] * uprev + β2end[i] * dt * k
-        if integrator.opts.adaptive
-            utilde = utilde + bhat2end[i] * dt * k
-        end
-    end
-
-    if integrator.opts.adaptive
-        atmp = calculate_residuals(
-            utilde, uprev, u, integrator.opts.abstol,
-            integrator.opts.reltol, integrator.opts.internalnorm, t
-        )
-        OrdinaryDiffEqCore.set_EEst!(integrator, integrator.opts.internalnorm(atmp, t))
-    end
-
-    integrator.u = u
+function perform_step!(integrator, cache::LowStorageRK3SpConstantCache, repeat_step = false)
+    return _perform_step_oop!(integrator, cache)
 end
 
 function initialize!(integrator, cache::LowStorageRK3SpCache)
-    (; k, fsalfirst) = cache
-
     integrator.kshortsize = 1
     resize!(integrator.k, integrator.kshortsize)
     return integrator.k[1] = integrator.fsalfirst
 end
 
-@muladd function perform_step!(integrator, cache::LowStorageRK3SpCache, repeat_step = false)
-    (; t, dt, uprev, u, f, p) = integrator
-    (; k, tmp, utilde, atmp, stage_limiter!, step_limiter!, thread) = cache
-    (; γ12end, γ22end, γ32end, δ2end, β1, β2end, c2end, bhat1, bhat2end) = cache.tab
-
-    # u1
-    f(integrator.fsalfirst, uprev, p, t)
-    OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
-    @.. broadcast = false thread = thread tmp = uprev
-    @.. broadcast = false thread = thread u = tmp + β1 * dt * integrator.fsalfirst
-    if integrator.opts.adaptive
-        @.. broadcast = false thread = thread utilde = bhat1 * dt * integrator.fsalfirst
-    end
-
-    # other stages
-    for i in eachindex(γ12end)
-        stage_limiter!(u, integrator, p, t + c2end[i] * dt)
-        f(k, u, p, t + c2end[i] * dt)
-        OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
-        @.. broadcast = false thread = thread tmp = tmp + δ2end[i] * u
-        @.. broadcast = false thread = thread u = γ12end[i] * u + γ22end[i] * tmp +
-            γ32end[i] * uprev + β2end[i] * dt * k
-        if integrator.opts.adaptive
-            @.. broadcast = false thread = thread utilde = utilde + bhat2end[i] * dt * k
-        end
-    end
-
-    stage_limiter!(u, integrator, p, t + dt)
-    step_limiter!(u, integrator, p, t + dt)
-
-    if integrator.opts.adaptive
-        calculate_residuals!(
-            atmp, utilde, uprev, u, integrator.opts.abstol,
-            integrator.opts.reltol, integrator.opts.internalnorm, t,
-            thread
-        )
-        OrdinaryDiffEqCore.set_EEst!(integrator, integrator.opts.internalnorm(atmp, t))
-    end
+function perform_step!(integrator, cache::LowStorageRK3SpCache, repeat_step = false)
+    return _perform_step_iip!(integrator, cache, cache.tab)
 end
 
-# 3S+ FSAL low storage methods: 3S methods adding another memory location for the embedded method (FSAL version)
 function initialize!(integrator, cache::LowStorageRK3SpFSALConstantCache)
-    integrator.fsalfirst = integrator.f(integrator.uprev, integrator.p, integrator.t) # Pre-start fsal
+    integrator.fsalfirst = integrator.f(integrator.uprev, integrator.p, integrator.t)
     OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
     integrator.kshortsize = 2
     integrator.k = typeof(integrator.k)(undef, integrator.kshortsize)
-
-    # Avoid undefined entries if k is an array of arrays
     integrator.fsallast = zero(integrator.fsalfirst)
     integrator.k[1] = integrator.fsalfirst
     return integrator.k[2] = integrator.fsallast
 end
 
-@muladd function perform_step!(
-        integrator, cache::LowStorageRK3SpFSALConstantCache,
-        repeat_step = false
+function perform_step!(
+        integrator, cache::LowStorageRK3SpFSALConstantCache, repeat_step = false
     )
-    (; t, dt, uprev, u, f, p) = integrator
-    (; γ12end, γ22end, γ32end, δ2end, β1, β2end, c2end, bhat1, bhat2end, bhatfsal) = cache
-
-    # u1
-    tmp = uprev
-    u = tmp + β1 * dt * integrator.fsalfirst
-    # Initialize utilde for JET
-    utilde = u
-    if integrator.opts.adaptive
-        utilde = bhat1 * dt * integrator.fsalfirst
-    end
-
-    # other stages
-    for i in eachindex(γ12end)
-        k = f(u, p, t + c2end[i] * dt)
-        OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
-        tmp = tmp + δ2end[i] * u
-        u = γ12end[i] * u + γ22end[i] * tmp + γ32end[i] * uprev + β2end[i] * dt * k
-        if integrator.opts.adaptive
-            utilde = utilde + bhat2end[i] * dt * k
-        end
-    end
-
-    # FSAL
-    integrator.fsallast = f(u, p, t + dt)
-    OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
-
-    if integrator.opts.adaptive
-        utilde = utilde + bhatfsal * dt * integrator.fsallast
-        atmp = calculate_residuals(
-            utilde, uprev, u, integrator.opts.abstol,
-            integrator.opts.reltol, integrator.opts.internalnorm, t
-        )
-        OrdinaryDiffEqCore.set_EEst!(integrator, integrator.opts.internalnorm(atmp, t))
-    end
-
-    integrator.k[1] = integrator.fsalfirst
-    integrator.k[2] = integrator.fsallast
-    integrator.u = u
+    return _perform_step_oop!(integrator, cache)
 end
 
 function initialize!(integrator, cache::LowStorageRK3SpFSALCache)
-    (; k, fsalfirst) = cache
-
     integrator.kshortsize = 2
     resize!(integrator.k, integrator.kshortsize)
     integrator.k[1] = integrator.fsalfirst
     integrator.k[2] = integrator.fsallast
-    integrator.f(integrator.fsalfirst, integrator.uprev, integrator.p, integrator.t) # FSAL for interpolation
+    integrator.f(integrator.fsalfirst, integrator.uprev, integrator.p, integrator.t)
     return OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
 end
 
-@muladd function perform_step!(
-        integrator, cache::LowStorageRK3SpFSALCache,
-        repeat_step = false
-    )
-    (; t, dt, uprev, u, f, p) = integrator
-    (; k, tmp, utilde, atmp, stage_limiter!, step_limiter!, thread) = cache
-    (;
-        γ12end, γ22end, γ32end, δ2end, β1, β2end,
-        c2end, bhat1, bhat2end, bhatfsal,
-    ) = cache.tab
-
-    # u1
-    @.. broadcast = false thread = thread tmp = uprev
-    @.. broadcast = false thread = thread u = tmp + β1 * dt * integrator.fsalfirst
-    if integrator.opts.adaptive
-        @.. broadcast = false thread = thread utilde = bhat1 * dt * integrator.fsalfirst
-    end
-
-    # other stages
-    for i in eachindex(γ12end)
-        stage_limiter!(u, integrator, p, t + c2end[i] * dt)
-        f(k, u, p, t + c2end[i] * dt)
-        OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
-        @.. broadcast = false thread = thread tmp = tmp + δ2end[i] * u
-        @.. broadcast = false thread = thread u = γ12end[i] * u + γ22end[i] * tmp +
-            γ32end[i] * uprev + β2end[i] * dt * k
-        if integrator.opts.adaptive
-            @.. broadcast = false thread = thread utilde = utilde + bhat2end[i] * dt * k
-        end
-    end
-
-    stage_limiter!(u, integrator, p, t + dt)
-    step_limiter!(u, integrator, p, t + dt)
-
-    # FSAL
-    f(k, u, p, t + dt)
-    OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
-
-    if integrator.opts.adaptive
-        @.. broadcast = false thread = thread utilde = utilde + bhatfsal * dt * k
-        calculate_residuals!(
-            atmp, utilde, uprev, u, integrator.opts.abstol,
-            integrator.opts.reltol, integrator.opts.internalnorm, t,
-            thread
-        )
-        OrdinaryDiffEqCore.set_EEst!(integrator, integrator.opts.internalnorm(atmp, t))
-    end
+function perform_step!(integrator, cache::LowStorageRK3SpFSALCache, repeat_step = false)
+    return _perform_step_iip!(integrator, cache, cache.tab)
 end
 
 # 2R+ low storage methods