Lowstorage 3s

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,139 @@
+@muladd function _perform_step_iip!(
+        integrator, cache, tab::LowStorageRKTableau{TwoN}
+    )
+    (; 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{TwoN}
+    )
+    (; 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
+
+@muladd function _perform_step_iip!(
+        integrator, cache, tab::LowStorageRKTableau{TwoC}
+    )
+    (; t, dt, u, f, p) = integrator
+    (; k, fsalfirst, tmp, stage_limiter!, step_limiter!, thread) = cache
+    (; A2end, B1, B2end, c2end) = tab
+
+    @.. broadcast = false thread = thread k = integrator.fsalfirst
+    @.. broadcast = false thread = thread u = u + B1 * dt * k
+
+    for i in eachindex(A2end)
+        @.. broadcast = false thread = thread tmp = u + A2end[i] * dt * k
+        f(k, tmp, p, t + c2end[i] * dt)
+        OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
+        @.. broadcast = false thread = thread u = u + B2end[i] * dt * k
+    end
+    step_limiter!(u, integrator, p, t + dt)
+    f(k, u, p, t + dt)
+    OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
+    return nothing
+end
+
+@muladd function _perform_step_oop!(
+        integrator, tab::LowStorageRKTableau{TwoC}
+    )
+    (; t, dt, u, f, p) = integrator
+    (; A2end, B1, B2end, c2end) = tab
+
+    k = integrator.fsalfirst = integrator.f(u, p, t)
+    integrator.k[1] = integrator.fsalfirst
+    OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
+    u = u + B1 * dt * k
+
+    for i in eachindex(A2end)
+        tmp = u + A2end[i] * dt * k
+        k = integrator.f(tmp, p, t + c2end[i] * dt)
+        OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
+        u = u + B2end[i] * dt * k
+    end
+
+    integrator.u = u
+    return nothing
+end
+
+@muladd function _perform_step_oop!(integrator, tab::LowStorageRK3SConstantCache)
+    (; t, dt, uprev, u, f, p) = integrator
+    (; γ12end, γ22end, γ32end, δ2end, β1, β2end, c2end) = tab
+
+    tmp = u
+    u = tmp + β1 * dt * integrator.fsalfirst
+
+    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
+    end
+
+    integrator.fsallast = f(u, p, t + dt)
+    OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
+    integrator.k[1] = integrator.fsalfirst
+    integrator.u = u
+    return nothing
+end
+
+@muladd function _perform_step_iip!(integrator, cache, tab::LowStorageRK3SConstantCache)
+    (; t, dt, uprev, u, f, p) = integrator
+    (; k, fsalfirst, tmp, stage_limiter!, step_limiter!, thread) = cache
+    (; γ12end, γ22end, γ32end, δ2end, β1, β2end, c2end) = tab
+
+    @.. broadcast = false thread = thread tmp = u
+    @.. broadcast = false thread = thread u = tmp + β1 * dt * integrator.fsalfirst
+
+    for i in eachindex(γ12end)
+        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
+    end
+
+    step_limiter!(u, integrator, p, t + dt)
+    f(k, u, p, t + dt)
+    OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
+    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{TwoN}(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{TwoN}(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{TwoN}(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{TwoN}(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{TwoN}(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{TwoN}(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{TwoN}(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{TwoN}(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{TwoN}(A2end, B1, B2end, c2end)
 end
 
 # 2C low storage methods introduced by Calvo, Franco, Rández (2004)
@@ -612,13 +605,6 @@ end
     thread::Thread
 end
 
-struct LowStorageRK2CConstantCache{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 CFRLDDRK64ConstantCache(::Type{T}, ::Type{T2}) where {T, T2}
     A2 = convert(T, 0.17985400977138)
     A3 = convert(T, 0.14081893152111)
@@ -642,7 +628,7 @@ function CFRLDDRK64ConstantCache(::Type{T}, ::Type{T2}) where {T, T2}
     c6 = convert(T2, 0.83050587987157)
     c2end = (c2, c3, c4, c5, c6)
 
-    return LowStorageRK2CConstantCache(A2end, B1, B2end, c2end)
+    return LowStorageRKTableau{TwoC}(A2end, B1, B2end, c2end)
 end
 
 function TSLDDRK74ConstantCache(::Type{T}, ::Type{T2}) where {T, T2}
@@ -671,7 +657,7 @@ function TSLDDRK74ConstantCache(::Type{T}, ::Type{T2}) where {T, T2}
     c7 = convert(T2, 0.91124223849547205)
     c2end = (c2, c3, c4, c5, c6, c7)
 
-    return LowStorageRK2CConstantCache(A2end, B1, B2end, c2end)
+    return LowStorageRKTableau{TwoC}(A2end, B1, B2end, c2end)
 end
 
 # 3S low storage methods introduced by Ketcheson