LowStorageRK: unify 2N family on LowStorageRKTableau

Author: Your Name

lib/OrdinaryDiffEqLowStorageRK/src/OrdinaryDiffEqLowStorageRK.jl

@@ -25,8 +25,10 @@ using Reexport
 include("arrayfuse.jl")
 include("algorithms.jl")
 include("alg_utils.jl")
+include("low_storage_tableaus.jl")
 include("low_storage_rk_caches.jl")
 include("low_storage_rk_perform_step.jl")
+include("generic_low_storage_perform_step.jl")
 
 import PrecompileTools
 import Preferences

lib/OrdinaryDiffEqLowStorageRK/src/generic_low_storage_perform_step.jl

@@ -0,0 +1,54 @@
+# Generic LowStorageRK perform_step! driven by `LowStorageRKTableau{form,...}`.
+# Plain dispatch on the form tag (no @generated); each form gets one method.
+
+@muladd function _perform_step_iip!(
+        integrator, cache, tab::LowStorageRKTableau{:two_n}
+    )
+    (; t, dt, u, f, p) = integrator
+    (; k, tmp, williamson_condition, stage_limiter!, step_limiter!, thread) = cache
+    (; A2end, B1, B2end, c2end) = tab
+
+    f(k, u, p, t)
+    OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
+    @.. broadcast=false thread=thread tmp = dt * k
+    @.. broadcast=false thread=thread u = u + B1 * tmp
+
+    for i in eachindex(A2end)
+        if williamson_condition
+            f(ArrayFuse(tmp, u, (A2end[i], dt, B2end[i])), u, p, t + c2end[i] * dt)
+        else
+            @.. broadcast=false thread=thread tmp = A2end[i] * tmp
+            stage_limiter!(u, integrator, p, t + c2end[i] * dt)
+            f(k, u, p, t + c2end[i] * dt)
+            @.. broadcast=false thread=thread tmp = tmp + dt * k
+            @.. broadcast=false thread=thread u = u + B2end[i] * tmp
+        end
+        OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
+    end
+    stage_limiter!(u, integrator, p, t + dt)
+    step_limiter!(u, integrator, p, t + dt)
+    return nothing
+end
+
+@muladd function _perform_step_oop!(
+        integrator, tab::LowStorageRKTableau{:two_n}
+    )
+    (; t, dt, u, f, p) = integrator
+    (; A2end, B1, B2end, c2end) = tab
+
+    tmp = dt * integrator.fsalfirst
+    u = u + B1 * tmp
+
+    for i in eachindex(A2end)
+        k = f(u, p, t + c2end[i] * dt)
+        OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
+        tmp = A2end[i] * tmp + dt * k
+        u = u + B2end[i] * tmp
+    end
+
+    integrator.k[1] = integrator.fsalfirst
+    integrator.fsalfirst = f(u, p, t + dt)
+    OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
+    integrator.u = u
+    return nothing
+end

lib/OrdinaryDiffEqLowStorageRK/src/low_storage_rk_caches.jl

@@ -17,13 +17,6 @@ get_fsalfirstlast(cache::LowStorageRKMutableCache, u) = (cache.fsalfirst, cache.
     thread::Thread
 end
 
-struct LowStorageRK2NConstantCache{N, T, T2} <: OrdinaryDiffEqConstantCache
-    A2end::NTuple{N, T} # A1 is always zero
-    B1::T
-    B2end::NTuple{N, T}
-    c2end::NTuple{N, T2} # c1 is always zero
-end
-
 function ORK256ConstantCache(::Type{T}, ::Type{T2}) where {T, T2}
     A2 = convert(T, -1.0)
     A3 = convert(T, -1.55798)
@@ -44,7 +37,7 @@ function ORK256ConstantCache(::Type{T}, ::Type{T2}) where {T, T2}
     c5 = convert(T2, 0.8)
     c2end = (c2, c3, c4, c5)
 
-    return LowStorageRK2NConstantCache(A2end, B1, B2end, c2end)
+    return LowStorageRKTableau{:two_n}(A2end, B1, B2end, c2end)
 end
 
 function alg_cache(
@@ -144,7 +137,7 @@ function CarpenterKennedy2N54ConstantCache(::Type{T}, ::Type{T2}) where {T, T2}
     c5 = convert(T2, 2802321613138 // 2924317926251)
     c2end = (c2, c3, c4, c5)
 
-    return LowStorageRK2NConstantCache(A2end, B1, B2end, c2end)
+    return LowStorageRKTableau{:two_n}(A2end, B1, B2end, c2end)
 end
 
 @cache mutable struct SHLDDRK_2NCache{
@@ -359,7 +352,7 @@ function SHLDDRK64ConstantCache(::Type{T}, ::Type{T2}) where {T, T2}
     c6 = convert(T2, 0.9271047)
     c2end = (c2, c3, c4, c5, c6)
 
-    return LowStorageRK2NConstantCache(A2end, B1, B2end, c2end)
+    return LowStorageRKTableau{:two_n}(A2end, B1, B2end, c2end)
 end
 
 function DGLDDRK73_CConstantCache(::Type{T}, ::Type{T2}) where {T, T2}
@@ -388,7 +381,7 @@ function DGLDDRK73_CConstantCache(::Type{T}, ::Type{T2}) where {T, T2}
     c7 = convert(T2, 0.998046608462379)
     c2end = (c2, c3, c4, c5, c6, c7)
 
-    return LowStorageRK2NConstantCache(A2end, B1, B2end, c2end)
+    return LowStorageRKTableau{:two_n}(A2end, B1, B2end, c2end)
 end
 
 function DGLDDRK84_CConstantCache(::Type{T}, ::Type{T2}) where {T, T2}
@@ -420,7 +413,7 @@ function DGLDDRK84_CConstantCache(::Type{T}, ::Type{T2}) where {T, T2}
     c8 = convert(T2, 0.919902896453866)
     c2end = (c2, c3, c4, c5, c6, c7, c8)
 
-    return LowStorageRK2NConstantCache(A2end, B1, B2end, c2end)
+    return LowStorageRKTableau{:two_n}(A2end, B1, B2end, c2end)
 end
 
 function DGLDDRK84_FConstantCache(::Type{T}, ::Type{T2}) where {T, T2}
@@ -452,7 +445,7 @@ function DGLDDRK84_FConstantCache(::Type{T}, ::Type{T2}) where {T, T2}
     c8 = convert(T2, 0.9484087623348481)
     c2end = (c2, c3, c4, c5, c6, c7, c8)
 
-    return LowStorageRK2NConstantCache(A2end, B1, B2end, c2end)
+    return LowStorageRKTableau{:two_n}(A2end, B1, B2end, c2end)
 end
 
 function NDBLSRK124ConstantCache(::Type{T}, ::Type{T2}) where {T, T2}
@@ -496,7 +489,7 @@ function NDBLSRK124ConstantCache(::Type{T}, ::Type{T2}) where {T, T2}
     c12 = convert(T2, 0.9032588871651854)
     c2end = (c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12)
 
-    return LowStorageRK2NConstantCache(A2end, B1, B2end, c2end)
+    return LowStorageRKTableau{:two_n}(A2end, B1, B2end, c2end)
 end
 
 function NDBLSRK134ConstantCache(::Type{T}, ::Type{T2}) where {T, T2}
@@ -543,7 +536,7 @@ function NDBLSRK134ConstantCache(::Type{T}, ::Type{T2}) where {T, T2}
     c13 = convert(T2, 0.9126827615920843)
     c2end = (c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13)
 
-    return LowStorageRK2NConstantCache(A2end, B1, B2end, c2end)
+    return LowStorageRKTableau{:two_n}(A2end, B1, B2end, c2end)
 end
 
 function NDBLSRK144ConstantCache(::Type{T}, ::Type{T2}) where {T, T2}
@@ -593,7 +586,7 @@ function NDBLSRK144ConstantCache(::Type{T}, ::Type{T2}) where {T, T2}
     c14 = convert(T2, 0.8734213127600976)
     c2end = (c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13, c14)
 
-    return LowStorageRK2NConstantCache(A2end, B1, B2end, c2end)
+    return LowStorageRKTableau{:two_n}(A2end, B1, B2end, c2end)
 end
 
 # 2C low storage methods introduced by Calvo, Franco, Rández (2004)

lib/OrdinaryDiffEqLowStorageRK/src/low_storage_rk_perform_step.jl

@@ -1,38 +1,15 @@
-# 2N low storage methods
-function initialize!(integrator, cache::LowStorageRK2NConstantCache)
-    integrator.fsalfirst = integrator.f(integrator.uprev, integrator.p, integrator.t) # Pre-start fsal
+# 2N low storage methods — dispatched via LowStorageRKTableau{:two_n}
+function initialize!(integrator, cache::LowStorageRKTableau{:two_n})
+    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::LowStorageRK2NConstantCache,
-        repeat_step = false
-    )
-    (; t, dt, u, f, p) = integrator
-    (; A2end, B1, B2end, c2end) = cache
-
-    # u1
-    tmp = dt * integrator.fsalfirst
-    u = u + B1 * tmp
-
-    # other stages
-    for i in eachindex(A2end)
-        k = f(u, p, t + c2end[i] * dt)
-        OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
-        tmp = A2end[i] * tmp + dt * k
-        u = u + B2end[i] * tmp
-    end
-
-    OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
-    integrator.k[1] = integrator.fsalfirst
-    integrator.fsalfirst = f(u, p, t + dt) # For interpolation, then FSAL'd
-    integrator.u = u
+function perform_step!(integrator, cache::LowStorageRKTableau{:two_n}, repeat_step = false)
+    return _perform_step_oop!(integrator, cache)
 end
 
 get_fsalfirstlast(cache::LowStorageRK2NCache, u) = (nothing, nothing)
@@ -42,35 +19,12 @@ function initialize!(integrator, cache::LowStorageRK2NCache)
     integrator.kshortsize = 1
     resize!(integrator.k, integrator.kshortsize)
     integrator.k[1] = k
-    integrator.f(k, integrator.uprev, integrator.p, integrator.t) # FSAL for interpolation
+    integrator.f(k, integrator.uprev, integrator.p, integrator.t)
     return OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
 end
 
-@muladd function perform_step!(integrator, cache::LowStorageRK2NCache, repeat_step = false)
-    (; t, dt, u, f, p) = integrator
-    (; k, tmp, williamson_condition, stage_limiter!, step_limiter!, thread) = cache
-    (; A2end, B1, B2end, c2end) = cache.tab
-
-    # u1
-    f(k, u, p, t)
-    OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
-    @.. broadcast = false thread = thread tmp = dt * k
-    @.. broadcast = false thread = thread u = u + B1 * tmp
-    # other stages
-    for i in eachindex(A2end)
-        if williamson_condition
-            f(ArrayFuse(tmp, u, (A2end[i], dt, B2end[i])), u, p, t + c2end[i] * dt)
-        else
-            @.. broadcast = false thread = thread tmp = A2end[i] * tmp
-            stage_limiter!(u, integrator, p, t + c2end[i] * dt)
-            f(k, u, p, t + c2end[i] * dt)
-            @.. broadcast = false thread = thread tmp = tmp + dt * k
-            @.. broadcast = false thread = thread u = u + B2end[i] * tmp
-        end
-        OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
-    end
-    stage_limiter!(u, integrator, p, t + dt)
-    step_limiter!(u, integrator, p, t + dt)
+function perform_step!(integrator, cache::LowStorageRK2NCache, repeat_step = false)
+    return _perform_step_iip!(integrator, cache, cache.tab)
 end
 
 # 2C low storage methods

lib/OrdinaryDiffEqLowStorageRK/src/low_storage_tableaus.jl

@@ -0,0 +1,24 @@
+# Unified tableau type for LowStorageRK families.
+#
+# Each method is identified by a `form` symbol tag carried as the first type
+# parameter. perform_step! dispatches on the tag via plain methods in
+# generic_low_storage_perform_step.jl.
+#
+# Currently supported forms:
+#   :two_n      — Williamson 2N      (A2end, B1, B2end, c2end)
+#
+# To add a new form: extend the field set if needed, add a constructor flavor,
+# and add the corresponding dispatch in generic_low_storage_perform_step.jl.
+
+struct LowStorageRKTableau{form, N, T, T2} <: OrdinaryDiffEqConstantCache
+    A2end::NTuple{N, T}
+    B1::T
+    B2end::NTuple{N, T}
+    c2end::NTuple{N, T2}
+end
+
+function LowStorageRKTableau{form}(
+        A2end::NTuple{N, T}, B1::T, B2end::NTuple{N, T}, c2end::NTuple{N, T2}
+    ) where {form, N, T, T2}
+    return LowStorageRKTableau{form, N, T, T2}(A2end, B1, B2end, c2end)
+end