Merge pull request #2566 from efaulhaber/drift-kick-drift

Implement the drift-kick-drift form of the Leapfrog method

Author: ChrisRackauckas

docs/src/dynamicalodeexplicit/SymplecticRK.md

@@ -47,6 +47,7 @@ sol = solve(prob, KahanLi8(), dt = 1 / 10)
 SymplecticEuler
 VelocityVerlet
 VerletLeapfrog
+LeapfrogDriftKickDrift
 PseudoVerletLeapfrog
 McAte2
 Ruth3

lib/OrdinaryDiffEqSymplecticRK/Project.toml

@@ -26,7 +26,7 @@ Random = "<0.0.1, 1"
 RecursiveArrayTools = "3.27.0"
 Reexport = "1.2.2"
 SafeTestsets = "0.1.0"
-Statistics = "1.11.1"
+Statistics = "<0.0.1, 1"
 Test = "<0.0.1, 1"
 julia = "1.10"
 

lib/OrdinaryDiffEqSymplecticRK/src/OrdinaryDiffEqSymplecticRK.jl

@@ -25,8 +25,8 @@ include("symplectic_caches.jl")
 include("symplectic_tableaus.jl")
 include("symplectic_perform_step.jl")
 
-export SymplecticEuler, VelocityVerlet, VerletLeapfrog, PseudoVerletLeapfrog,
-       McAte2, Ruth3, McAte3, CandyRoz4, McAte4, McAte42, McAte5,
+export SymplecticEuler, VelocityVerlet, VerletLeapfrog, LeapfrogDriftKickDrift,
+       PseudoVerletLeapfrog, McAte2, Ruth3, McAte3, CandyRoz4, McAte4, McAte42, McAte5,
        CalvoSanz4, Yoshida6, KahanLi6, McAte8, KahanLi8, SofSpa10
 
 end

lib/OrdinaryDiffEqSymplecticRK/src/alg_utils.jl

@@ -1,6 +1,7 @@
 alg_order(alg::SymplecticEuler) = 1
 alg_order(alg::VelocityVerlet) = 2
 alg_order(alg::VerletLeapfrog) = 2
+alg_order(alg::LeapfrogDriftKickDrift) = 2
 alg_order(alg::PseudoVerletLeapfrog) = 2
 alg_order(alg::McAte2) = 2
 alg_order(alg::Ruth3) = 3

lib/OrdinaryDiffEqSymplecticRK/src/algorithms.jl

@@ -24,14 +24,38 @@ publisher={APS}
     verlet1967, "", "")
 struct VelocityVerlet <: OrdinaryDiffEqPartitionedAlgorithm end
 
-@doc generic_solver_docstring("2nd order explicit symplectic integrator.",
+monaghan2005 = """
+@article{monaghan2005,
+	title = {Smoothed particle hydrodynamics},
+	author = {Monaghan, Joseph J.},
+	year = {2005},
+	journal = {Reports on Progress in Physics},
+	volume = {68},
+	number = {8},
+	pages = {1703--1759},
+	doi = {10.1088/0034-4885/68/8/R01},
+}
+"""
+
+@doc generic_solver_docstring(
+    "2nd order explicit symplectic integrator. Kick-drift-kick form. Requires only one evaluation of `f1` per step.",
     "VerletLeapfrog",
     "Symplectic Runge-Kutta Methods",
-    verlet1967, "", "")
+    monaghan2005, "", "")
 struct VerletLeapfrog <: OrdinaryDiffEqPartitionedAlgorithm end
 
 OrdinaryDiffEqCore.default_linear_interpolation(alg::VerletLeapfrog, prob) = true
 
+@doc generic_solver_docstring(
+    "2nd order explicit symplectic integrator. Drift-kick-drift form of `VerletLeapfrog`
+designed to work when `f1` depends on `v`. Requires two evaluation of `f1` per step.",
+    "LeapfrogDriftKickDrift",
+    "Symplectic Runge-Kutta Methods",
+    monaghan2005, "", "")
+struct LeapfrogDriftKickDrift <: OrdinaryDiffEqPartitionedAlgorithm end
+
+OrdinaryDiffEqCore.default_linear_interpolation(alg::LeapfrogDriftKickDrift, prob) = true
+
 @doc generic_solver_docstring("2nd order explicit symplectic integrator.",
     "PseudoVerletLeapfrog",
     "Symplectic Runge-Kutta Methods",

lib/OrdinaryDiffEqSymplecticRK/src/symplectic_caches.jl

@@ -57,6 +57,38 @@ function alg_cache(alg::VelocityVerlet, u, rate_prototype, ::Type{uEltypeNoUnits
     VelocityVerletConstantCache(uEltypeNoUnits(1 // 2))
 end
 
+@cache struct LeapfrogDriftKickDriftCache{uType, rateType, uEltypeNoUnits} <:
+              OrdinaryDiffEqMutableCache
+    u::uType
+    uprev::uType
+    tmp::uType
+    k::rateType
+    fsalfirst::rateType
+    half::uEltypeNoUnits
+end
+
+struct LeapfrogDriftKickDriftConstantCache{uEltypeNoUnits} <: HamiltonConstantCache
+    half::uEltypeNoUnits
+end
+
+function alg_cache(alg::LeapfrogDriftKickDrift, u, rate_prototype, ::Type{uEltypeNoUnits},
+        ::Type{uBottomEltypeNoUnits}, ::Type{tTypeNoUnits}, uprev, uprev2, f, t,
+        dt, reltol, p, calck,
+        ::Val{true}) where {uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits}
+    tmp = zero(rate_prototype)
+    k = zero(rate_prototype)
+    fsalfirst = zero(rate_prototype)
+    half = uEltypeNoUnits(1 // 2)
+    LeapfrogDriftKickDriftCache(u, uprev, k, tmp, fsalfirst, half)
+end
+
+function alg_cache(alg::LeapfrogDriftKickDrift, u, rate_prototype, ::Type{uEltypeNoUnits},
+        ::Type{uBottomEltypeNoUnits}, ::Type{tTypeNoUnits}, uprev, uprev2, f, t,
+        dt, reltol, p, calck,
+        ::Val{false}) where {uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits}
+    LeapfrogDriftKickDriftConstantCache(uEltypeNoUnits(1 // 2))
+end
+
 @cache struct VerletLeapfrogCache{uType, rateType, uEltypeNoUnits} <:
               OrdinaryDiffEqMutableCache
     u::uType
@@ -436,6 +468,6 @@ end
 
 function get_fsalfirstlast(
         cache::Union{HamiltonMutableCache, VelocityVerletCache, VerletLeapfrogCache,
-            SymplecticEulerCache}, u)
+            SymplecticEulerCache, LeapfrogDriftKickDriftCache}, u)
     (cache.fsalfirst, cache.k)
 end

lib/OrdinaryDiffEqSymplecticRK/src/symplectic_perform_step.jl

@@ -78,11 +78,13 @@ end
 # If called with different functions (which are possible in the Hamiltonian case)
 # an exception is thrown to avoid silently calculate wrong results.
 function verify_f2(f, p, q, pa, t, ::Any,
-        ::C) where {C <: Union{HamiltonConstantCache, VerletLeapfrogConstantCache}}
+        ::C) where {C <: Union{HamiltonConstantCache, VerletLeapfrogConstantCache,
+        LeapfrogDriftKickDriftConstantCache}}
     f(p, q, pa, t)
 end
 function verify_f2(f, res, p, q, pa, t, ::Any,
-        ::C) where {C <: Union{HamiltonMutableCache, VerletLeapfrogCache}}
+        ::C) where {C <: Union{HamiltonMutableCache, VerletLeapfrogCache,
+        LeapfrogDriftKickDriftCache}}
     f(res, p, q, pa, t)
 end
 
@@ -128,8 +130,8 @@ function store_symp_state!(integrator, ::OrdinaryDiffEqMutableCache, kdu, ku)
 end
 
 function initialize!(integrator,
-        cache::C) where {C <: Union{
-        HamiltonMutableCache, VelocityVerletCache, VerletLeapfrogCache}}
+        cache::C) where {C <: Union{HamiltonMutableCache, VelocityVerletCache,
+        VerletLeapfrogCache, LeapfrogDriftKickDriftCache}}
     integrator.kshortsize = 2
     resize!(integrator.k, integrator.kshortsize)
     integrator.k[1] = integrator.fsalfirst
@@ -144,8 +146,8 @@ function initialize!(integrator,
 end
 
 function initialize!(integrator,
-        cache::C) where {C <: Union{
-        HamiltonConstantCache, VelocityVerletConstantCache, VerletLeapfrogConstantCache}}
+        cache::C) where {C <: Union{HamiltonConstantCache, VelocityVerletConstantCache,
+        VerletLeapfrogConstantCache, LeapfrogDriftKickDriftConstantCache}}
     integrator.kshortsize = 2
     integrator.k = typeof(integrator.k)(undef, integrator.kshortsize)
 
@@ -207,13 +209,11 @@ end
     du = duprev + dt * half * kduprev
 
     # update position
-    tnew = t + half * dt
-    ku = f.f2(du, uprev, p, tnew)
+    ku = f.f2(du, uprev, p, t + half * dt)
     u = uprev + dt * ku
 
     # update velocity
-    tnew = tnew + half * dt
-    kdu = f.f1(du, u, p, tnew)
+    kdu = f.f1(du, u, p, t + dt)
     du = du + dt * half * kdu
 
     OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
@@ -226,26 +226,83 @@ end
     duprev, uprev, kduprev, _ = load_symp_state(integrator)
     du, u, kdu, ku = alloc_symp_state(integrator)
 
-    # Kick-Drift-Kick scheme of the Verlet Leapfrog method:
+    # kick-drift-kick scheme of the Leapfrog method:
     # update velocity
     half = cache.half
     @.. broadcast=false du=duprev + dt * half * kduprev
 
     # update position
-    tnew = t + half * dt
-    f.f2(ku, du, uprev, p, tnew)
+    f.f2(ku, du, uprev, p, t + half * dt)
     @.. broadcast=false u=uprev + dt * ku
 
     # update velocity
-    tnew = tnew + half * dt
-    f.f1(kdu, du, u, p, tnew)
+    f.f1(kdu, du, u, p, t + dt)
     @.. broadcast=false du=du + dt * half * kdu
 
     OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
     integrator.stats.nf2 += 1
     store_symp_state!(integrator, cache, kdu, ku)
 end
 
+@muladd function perform_step!(integrator, cache::LeapfrogDriftKickDriftConstantCache,
+        repeat_step = false)
+    @unpack t, dt, f, p = integrator
+    duprev, uprev, _, _ = load_symp_state(integrator)
+
+    # drift-kick-drift scheme of the Leapfrog method, allowing for f1 to depend on v:
+    # update position half step
+    half = cache.half
+    ku = f.f2(duprev, uprev, p, t)
+    u = uprev + dt * half * ku
+
+    # update velocity half step
+    kdu = f.f1(duprev, uprev, p, t)
+    du = duprev + dt * half * kdu
+
+    # update velocity (add to previous full step velocity)
+    # note that this extra step is only necessary if f1 depends on v/du (or t)
+    kdu = f.f1(du, u, p, t + half * dt)
+    du = duprev + dt * kdu
+
+    # update position (add to half step position)
+    ku = f.f2(du, u, p, t + dt)
+    u = u + dt * half * ku
+
+    OrdinaryDiffEqCore.increment_nf!(integrator.stats, 2)
+    integrator.stats.nf2 += 2
+    store_symp_state!(integrator, cache, du, u, kdu, ku)
+end
+
+@muladd function perform_step!(integrator, cache::LeapfrogDriftKickDriftCache,
+        repeat_step = false)
+    @unpack t, dt, f, p = integrator
+    duprev, uprev, _, _ = load_symp_state(integrator)
+    du, u, kdu, ku = alloc_symp_state(integrator)
+
+    # drift-kick-drift scheme of the Leapfrog method, allowing for f1 to depend on v:
+    # update position half step
+    half = cache.half
+    f.f2(ku, duprev, uprev, p, t)
+    @.. broadcast=false u=uprev + dt * half * ku
+
+    # update velocity half step
+    f.f1(kdu, duprev, uprev, p, t)
+    @.. broadcast=false du=duprev + dt * half * kdu
+
+    # update velocity (add to previous full step velocity)
+    # note that this extra step is only necessary if f1 depends on v/du (or t)
+    f.f1(kdu, du, u, p, t + half * dt)
+    @.. broadcast=false du=duprev + dt * kdu
+
+    # update position (add to half step position)
+    f.f2(ku, du, u, p, t + dt)
+    @.. broadcast=false u=u + dt * half * ku
+
+    OrdinaryDiffEqCore.increment_nf!(integrator.stats, 2)
+    integrator.stats.nf2 += 2
+    store_symp_state!(integrator, cache, kdu, ku)
+end
+
 @muladd function perform_step!(integrator, cache::Symplectic2ConstantCache,
         repeat_step = false)
     @unpack t, dt, f, p = integrator

lib/OrdinaryDiffEqSymplecticRK/test/symplectic_convergence.jl

@@ -62,6 +62,9 @@ position_error = :final => [mean(sim[i].u[2].x[1] - sim[i].u_analytic[2].x[1])
 sim = test_convergence(dts, prob, VerletLeapfrog(), dense_errors = true)
 @test sim.𝒪est[:l2]≈2 rtol=1e-1
 @test sim.𝒪est[:L2]≈2 rtol=1e-1
+sim = test_convergence(dts, prob, LeapfrogDriftKickDrift(), dense_errors = true)
+@test sim.𝒪est[:l2]≈2 rtol=1e-1
+@test sim.𝒪est[:L2]≈2 rtol=1e-1
 sim = test_convergence(dts, prob, PseudoVerletLeapfrog(), dense_errors = true)
 @test sim.𝒪est[:l2]≈2 rtol=1e-1
 @test sim.𝒪est[:L2]≈2 rtol=1e-1
@@ -151,6 +154,9 @@ position_error = :final => [mean(sim[i].u[2].x[1] - sim[i].u_analytic[2].x[1])
 sim = test_convergence(dts, prob, VerletLeapfrog(), dense_errors = true)
 @test sim.𝒪est[:l2]≈2 rtol=1e-1
 @test sim.𝒪est[:L2]≈2 rtol=1e-1
+sim = test_convergence(dts, prob, LeapfrogDriftKickDrift(), dense_errors = true)
+@test sim.𝒪est[:l2]≈2 rtol=1e-1
+@test sim.𝒪est[:L2]≈2 rtol=1e-1
 sim = test_convergence(dts, prob, PseudoVerletLeapfrog(), dense_errors = true)
 @test sim.𝒪est[:l2]≈2 rtol=1e-1
 @test sim.𝒪est[:L2]≈2 rtol=1e-1
@@ -202,3 +208,28 @@ dts = 1.0 ./ 2.0 .^ (2:-1:-2)
 sim = test_convergence(dts, prob, SofSpa10(), dense_errors = true)
 @test sim.𝒪est[:l2]≈10 rtol=1e-1
 @test sim.𝒪est[:L2]≈4 rtol=1e-1
+
+################# f1 dependent on v
+
+println("f1 dependent on v")
+
+u0 = fill(0.0, 2)
+v0 = ones(2)
+function f1_v(dv, v, u, p, t)
+    dv .= v
+end
+function f2_v(du, v, u, p, t)
+    du .= v
+end
+function f_v_analytic(y0, p, x)
+    v0, u0 = y0.x
+    ArrayPartition(v0 * exp(x), v0 * exp(x) - v0 + u0)
+end
+ff_v = DynamicalODEFunction(f1_v, f2_v; analytic = f_v_analytic)
+prob = DynamicalODEProblem(ff_v, v0, u0, (0.0, 5.0))
+
+dts = 1 .// 2 .^ (6:-1:3)
+# LeapfrogDriftKickDrift
+sim = test_convergence(dts, prob, LeapfrogDriftKickDrift(), dense_errors = true)
+@test sim.𝒪est[:l2]≈2 rtol=1e-1
+@test sim.𝒪est[:L2]≈2 rtol=1e-1

lib/OrdinaryDiffEqSymplecticRK/test/symplectic_tests.jl

@@ -7,6 +7,7 @@ using OrdinaryDiffEqRKN
 const ALGOS = ((SymplecticEuler, true, 1),
     (VelocityVerlet, false, 2),
     (VerletLeapfrog, true, 2),
+    (LeapfrogDriftKickDrift, true, 2),
     (PseudoVerletLeapfrog, true, 2),
     (McAte2, true, 2),
     (Ruth3, true, 3),

src/OrdinaryDiffEq.jl

@@ -113,8 +113,8 @@ using OrdinaryDiffEqFeagin
 export Feagin10, Feagin12, Feagin14
 
 using OrdinaryDiffEqSymplecticRK
-export SymplecticEuler, VelocityVerlet, VerletLeapfrog, PseudoVerletLeapfrog,
-       McAte2, Ruth3, McAte3, CandyRoz4, McAte4, McAte42, McAte5,
+export SymplecticEuler, VelocityVerlet, VerletLeapfrog, LeapfrogDriftKickDrift,
+       PseudoVerletLeapfrog, McAte2, Ruth3, McAte3, CandyRoz4, McAte4, McAte42, McAte5,
        CalvoSanz4, Yoshida6, KahanLi6, McAte8, KahanLi8, SofSpa10
 
 using OrdinaryDiffEqRKN