From 5ea1d65b2d25f79d499c3df6fc32e5969e56efb2 Mon Sep 17 00:00:00 2001
From: Kipton Barros <kbarros@gmail.com>
Date: Mon, 8 Jan 2024 10:25:13 -0700
Subject: [PATCH] Add `check_timestep`

This should be used after a time-integration to verify expected accuracy.
---
 docs/src/versions.md       |   2 +
 examples/03_LLD_CoRh2O4.jl |   5 ++
 examples/04_GSD_FeI2.jl    |   7 +++
 src/Integrators.jl         | 120 +++++++++++++++++++++++++++++++++++--
 src/OnlineStatistics.jl    |  27 +++++++++
 src/Sunny.jl               |   4 +-
 6 files changed, 158 insertions(+), 7 deletions(-)
 create mode 100644 src/OnlineStatistics.jl

diff --git a/docs/src/versions.md b/docs/src/versions.md
index e570f6576..db62a3061 100644
--- a/docs/src/versions.md
+++ b/docs/src/versions.md
@@ -3,6 +3,8 @@
 ## v0.5.9
 (In development)
 
+* New function [`check_timestep`](@ref) to verify accuracy of spin dynamics
+  integration.
 
 ## v0.5.8
 (Jan 4, 2024)
diff --git a/examples/03_LLD_CoRh2O4.jl b/examples/03_LLD_CoRh2O4.jl
index 373f03181..b7e3054eb 100644
--- a/examples/03_LLD_CoRh2O4.jl
+++ b/examples/03_LLD_CoRh2O4.jl
@@ -50,6 +50,11 @@ for _ in 1:1000
 end
 plot(energies, color=:blue, figure=(size=(600,300),), axis=(xlabel="Time steps", ylabel="Energy (meV)"))
 
+# After running a Langevin trajectory, it is a good practice to call
+# [`check_timestep`](@ref).
+
+check_timestep(langevin; tol=1e-2)
+
 # Thermal fluctuations are apparent in the spin configuration.
 
 S_ref = sys.dipoles[1,1,1,1]
diff --git a/examples/04_GSD_FeI2.jl b/examples/04_GSD_FeI2.jl
index 09c50cfd6..117324520 100644
--- a/examples/04_GSD_FeI2.jl
+++ b/examples/04_GSD_FeI2.jl
@@ -102,6 +102,11 @@ for _ in 1:20_000
     step!(sys, langevin)
 end
 
+# After running a Langevin trajectory, it is a good practice to call
+# [`check_timestep`](@ref).
+
+check_timestep(langevin; tol=1e-2)
+
 # Although thermal fluctuations are present, the correct antiferromagnetic order
 # (2 up, 2 down) is apparent.
 
@@ -131,6 +136,8 @@ for _ in 1:10_000
     step!(sys_large, langevin)
 end
 
+check_timestep(langevin; tol=1e-2)
+
 # The next step is to collect correlation data ``S^{\alpha\beta}``. This will
 # involve sampling spin configurations from thermal equilibrium, and then
 # integrating [an energy-conserving generalized classical spin
diff --git a/src/Integrators.jl b/src/Integrators.jl
index 88283bc5f..9b0128378 100644
--- a/src/Integrators.jl
+++ b/src/Integrators.jl
@@ -53,10 +53,14 @@ mutable struct Langevin
     λ   :: Float64
     kT  :: Float64
 
+    # Averages the square of the dynamical drift term for use in
+    # `check_timestep`.
+    drift_squared :: OnlineStatistics{Float64}
+
     function Langevin(Δt; λ, kT)
         Δt <= 0 && error("Select positive Δt")
-        return new(Δt, λ, kT)
-    end    
+        return new(Δt, λ, kT, OnlineStatistics{Float64}())
+    end
 end
 
 """
@@ -80,13 +84,95 @@ mutable struct ImplicitMidpoint
     Δt   :: Float64
     atol :: Float64
 
+    # Averages the square of the dynamical drift term for use in
+    # `check_timestep`.
+    drift_squared :: OnlineStatistics{Float64}
+
     function ImplicitMidpoint(Δt; atol=1e-12)
         Δt <= 0 && error("Select positive Δt")
-        return new(Δt, atol)
+        return new(Δt, atol, OnlineStatistics{Float64}())
     end    
 end
 
 
+function reset_statistics(integrator::Union{Langevin, ImplicitMidpoint})
+    integrator.drift_squared = OnlineStatistics{Float64}()
+end
+
+"""
+    check_timestep(integrator; tol)
+
+Prints a suggested timestep scale for the relative error tolerance `tol`, based
+on statistics collected from previous time integration. A reasonable tolerance
+might be of order `tol = 1e-2`, but should be selected on a case-by-case basis.
+Note that `tol` is interpreted as the error for the deterministic part of a
+single integration timestep, which scales like `Δt²`. Because the stochastic
+Heun integration scheme is weakly convergent of order-1, errors in the estimates
+of averaged observables may instead scale like ` √tol ~ Δt`. It is the
+responsibility of the user to choose `tol` accordingly.
+
+The suggested timestep will be inversely proportional to the largest energy
+scale that plays a role in the dynamics. If there is a strong Langvin coupling
+``λ`` to the thermal bath, it will effectively rescale the energy scale.
+
+Timestep statistics stored in `integrator` will be reset after calling this
+function.
+"""
+function check_timestep(langevin::Langevin; tol, reset=true)
+    (; Δt, λ, kT) = langevin
+
+    drift_magnitude = sqrt(mean(langevin.drift_squared))
+    if isnan(drift_magnitude)
+        error("Run trajectory first to collect statistics")
+    end
+
+    # Heun integration is second-order accurate, so the local error from each
+    # deterministic timestep scales as dθ². Angular displacement per timestep dθ
+    # scales like dt |drift|, yielding err1 ~ (dt |drift|)^2
+    #
+    # The "relevant" error introduced by thermal noise should be quantified in
+    # terms of the effect on statistical observables. To avoid subtleties in
+    # defining this error, we instead naïvely assume it continues to be second
+    # order in `dt`. To determine the proportionality constant, consider the
+    # high-T limit, where each spin undergoes Brownian motion. Here, the
+    # diffusion constant D ~ λ kT sets an inverse time-scale. This implies err2
+    # ~ (dt λ kT)².
+    #
+    # The total error (err1 + err2) should be less than the target tolerance.
+    # After some algebra, this implies,
+    #
+    # dt ≲ sqrt(tol / (c₁² |drift|² + c₂² λ² kT²))
+    #
+    # for some empirical constants c₁ and c₂.
+
+    c1 = 1.0
+    c2 = 1.0
+    Δt_bound = sqrt(tol / ((c1*drift_magnitude)^2 + (c2*λ*kT)^2))
+
+    if Δt_bound/2 < Δt < 2Δt_bound
+        opinion = "reasonable."
+    elseif Δt_bound/5 < Δt < 5Δt_bound
+        opinion = "_marginal_."
+    elseif Δt <= Δt_bound/5
+        opinion = "SMALL!"
+    elseif Δt >= 5Δt_bound
+        opinion = "LARGE!"
+    end
+
+    λkTstr, tolstr, Δtstr, boundstr = number_to_simple_string.((λ*kT, tol, Δt, Δt_bound); digits=4)
+
+    println("Suggest Δt ~ $boundstr for tol = $tolstr.")
+    println("Current Δt = $Δtstr is $opinion")
+    if c2*λ*kT > c1*drift_magnitude
+        println("Thermal noise contributes significantly, with λ*kT = $λkTstr.")
+    end
+
+    reset && reset_statistics(langevin)
+
+    return
+end
+
+
 ################################################################################
 # Dipole integration
 ################################################################################
@@ -122,7 +208,19 @@ function step!(sys::System{0}, integrator::Langevin)
     set_energy_grad_dipoles!(∇E, s₁, sys)
     @. s = s + 0.5 * Δt * (f₁ + rhs_dipole(s₁, -∇E, λ)) + 0.5 * √Δt * (r₁ + rhs_dipole(s₁, ξ))
     @. s = normalize_dipole(s, sys.κs)
-    nothing
+
+    # Collect statistics about the magnitude of the drift term squared. For the
+    # energy-conserving part, it is important to use |∇E|² rather than |s×∇E|²,
+    # because the former gives direct information about the frequency of
+    # oscillation, while the latter may artificially vanish (consider, e.g., a
+    # spin nearly aligned with an external field).
+    for i in eachsite(sys)
+        s0 = sys.κs[i] # == norm(s[i])
+        ∇E² = ∇E[i]' * ∇E[i]
+        accum!(integrator.drift_squared, (1 + (s0*λ)^2) * ∇E²)
+    end
+
+    return
 end
 
 # The spherical midpoint method, Phys. Rev. E 89, 061301(R) (2014)
@@ -178,8 +276,9 @@ end
 ################################################################################
 # SU(N) integration
 ################################################################################
-@inline function proj(a::T, Z::T)  where T <: CVec
-    (a - ((Z' * a) * Z))  
+
+@inline function proj(a::T, Z::T) where T <: CVec
+    a - Z * ((Z' * a) / (Z' * Z))
 end
 
 function step!(sys::System{N}, integrator::Langevin) where N
@@ -200,6 +299,15 @@ function step!(sys::System{N}, integrator::Langevin) where N
 
     # Coordinate dipole data
     @. sys.dipoles = expected_spin(Z)
+
+    # Collect statistics about the magnitude of the drift term squared. Here,
+    # |HZ|² plays the role of |∇E|² in the dipole case.
+    for i in eachsite(sys)
+        HZ² = HZ[i]' * HZ[i]
+        accum!(integrator.drift_squared, (1 + integrator.λ^2) * HZ²)
+    end
+
+    return
 end
 
 function rhs_langevin!(ΔZ::Array{CVec{N}, 4}, Z::Array{CVec{N}, 4}, ξ::Array{CVec{N}, 4},
diff --git a/src/OnlineStatistics.jl b/src/OnlineStatistics.jl
new file mode 100644
index 000000000..b56cfae65
--- /dev/null
+++ b/src/OnlineStatistics.jl
@@ -0,0 +1,27 @@
+
+mutable struct OnlineStatistics{T}
+    n::Int
+    μ::T
+    M::Float64
+ 
+    function OnlineStatistics{T}() where T
+        new(0, T(NaN), T(NaN))
+    end
+end
+
+function accum!(o::OnlineStatistics{T}, x) where T
+    o.n += 1
+    if o.n == 1
+        o.μ = x
+        o.M = 0.0
+    else
+        μ_prev = o.μ
+        o.μ += (x - μ_prev) / o.n
+        o.M += real((x - o.μ)' * (x - μ_prev))
+    end
+    o
+end
+
+# If Statistics becomes in scope, define methods Statistics.{mean, var} instead
+mean(o::OnlineStatistics) = o.μ
+var(o::OnlineStatistics; corrected=true) = corrected ? o.M/(o.n-1) : o.M/o.n
diff --git a/src/Sunny.jl b/src/Sunny.jl
index a56e21ef1..624de97ed 100644
--- a/src/Sunny.jl
+++ b/src/Sunny.jl
@@ -39,6 +39,8 @@ const HermitianC64 = Hermitian{ComplexF64, Matrix{ComplexF64}}
     end
 end
 
+include("OnlineStatistics.jl")
+
 include("Operators/Spin.jl")
 include("Operators/Rotation.jl")
 include("Operators/Stevens.jl")
@@ -82,7 +84,7 @@ include("Reshaping.jl")
 export reshape_supercell, resize_supercell, repeat_periodically
 
 include("Integrators.jl")
-export Langevin, ImplicitMidpoint, step!
+export Langevin, ImplicitMidpoint, step!, check_timestep
 
 include("Optimization.jl")
 export minimize_energy!