Hokseson

src/DynamicsMethods/IntegrationAlgorithms/mdef_baoab.jl

@@ -9,6 +9,18 @@
 
 StochasticDiffEq.alg_compatible(::DiffEqBase.AbstractSDEProblem,::MDEF_BAOAB) = true
 
+"""
+    MDEF_BAOAB():
+    We define two types of variables structure in the cache:
+    - `MDEF_BAOABConstantCache` is used for constant variables, which are not changed during the integration.
+
+    - `MDEF_BAOABCache` is used for variables that are changed during the integration.
+
+    Generally, the mutableChahe is faster than the constant cache. 
+
+    Mathmatically speaking, they are doing the same things. MutableCache is more well-optimize for the Julia compiler.
+"""
+
 struct MDEF_BAOABConstantCache{uType,uEltypeNoUnits} <: StochasticDiffEq.StochasticDiffEqConstantCache
     k::uType
     half::uEltypeNoUnits
@@ -28,6 +40,9 @@
     c2::Matrix{uEltypeNoUnits}
 end
 
+"""
+    Insecting the inputs into a Cache structure.
+"""
 function StochasticDiffEq.alg_cache(::MDEF_BAOAB,prob,u,ΔW,ΔZ,p,rate_prototype,noise_rate_prototype,jump_rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,f,t,dt,::Type{Val{false}})
     k = zero(rate_prototype.x[1])
     MDEF_BAOABConstantCache(k, uEltypeNoUnits(1//2))
@@ -52,6 +67,15 @@
     MDEF_BAOABCache(tmp, utmp, dutmp, k, flatdutmp, tmp1, tmp2, gtmp, noise, half, c1, c2)
 end
 
+"""
+    Initialize the the velocities du and positions u.
+
+    k : gradient of the potential energy
+
+    integrator.f.f1 : the gradient operator ---> acceleration!() in diabtic_mdef.jl
+
+"""
+
 function StochasticDiffEq.initialize!(integrator, cache::MDEF_BAOABConstantCache)
     @unpack t,uprev,p = integrator
     du1 = uprev.x[1]
@@ -68,6 +92,29 @@
     integrator.f.f1(cache.k,du1,u1,p,t)
 end
 
+"""
+    perform_step! perform a single step of BAOAB algorithm.
+
+    uprev: previous configureation(velocities and positions)
+    dt : length of time step
+    t : current time
+    k : gradient of the potential energy in terms of velocities
+    p : parameters like temperature, mass, etc.
+
+    The B and A steps are the same as Leimkuhler and Matthews (2013) and the O step is the same as Bussi and Parrinello (2007).
+
+    B and A step: https://academic.oup.com/amrx/article/2013/1/34/166771
+
+    O step : general part https://academic.oup.com/amrx/article/2013/1/34/166771
+
+             friction part https://journals.aps.org/pre/abstract/10.1103/PhysRevE.75.056707 
+                           and https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.118.256001 (supplementary material)
+
+    friction tensor Λ: https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.118.256001 (supplementary material)
+
+
+"""
+
 @muladd function StochasticDiffEq.perform_step!(integrator,cache::MDEF_BAOABConstantCache,f=integrator.f)
     @unpack t,dt,sqdt,uprev,p,W = integrator
     @unpack k, half = cache
@@ -81,14 +128,16 @@
     u2 = u1 + half*dt*du2
 
     # O
-    Λ = integrator.g(u2,p,t+dt*half)
+    Λ = integrator.g(u2,p,t+dt*half) # friction tensor
+    # noise strength: σ = square root of (temperature / mass) for each atom 
     σ = sqrt(get_temperature(p, t+dt*half)) ./ sqrt.(repeat(p.atoms.masses; inner=ndofs(p)))
+    # eigen decomposition of Λ
     γ, c = LAPACK.syev!('V', 'U', Λ) # symmetric eigen
     clamp!(γ, 0, Inf)
     c1 = diagm(exp.(-γ.*dt))
     c2 = diagm(sqrt.(1 .- diag(c1).^2))
-    noise = σ.*W.dW[:] / sqdt
-    du3 = c*c1*c'*du2[:] + c*c2*c'*noise
+    noise = σ.*W.dW[:] / sqdt # W.dW[:] noise rescaled by square root of time step to make it normal distributed
+    du3 = c*c1*c'*du2[:] + c*c2*c'*noise # O step from Leimkuhler and Matthews (2013) 
     du3 = reshape(du3, size(du2))
 
     # A
@@ -101,6 +150,12 @@
     integrator.u = ArrayPartition((du, u))
 end
 
+"""
+    The details of the following perform_step! using mutableChahe can be found in steps.jl in intergrationAlgorithm.
+
+    Before the checking, you are recommonded to read the function with ConstantCache input first.
+"""
+
 @muladd function StochasticDiffEq.perform_step!(integrator,cache::MDEF_BAOABCache,f=integrator.f)
     @unpack t,dt,sqdt,uprev,u,p,W = integrator
     @unpack utmp, dutmp, k, half, gtmp = cache

test/Dynamics/mdef_baoab.jl

@@ -6,6 +6,7 @@ using DiffEqNoiseProcess
 using DiffEqDevTools
 using Random
 using Logging
+# Set up the seed 
 Random.seed!(100)
 
 u0 = zeros(2,2)