Add diffusion Monte Carlo

docs/make.jl

@@ -23,6 +23,7 @@ pages = gem_only ?
         "Quantics Tensor Train" => "QTT.md",
         "Variational Neural Network" => "VNN.md",
         "Variational Monte Carlo" => "VMC.md",
+        "Diffusion Monte Carlo" => "DMC.md",
         "Developer Guide" => "developer.md",
         "API reference" => "API.md",
     ]

docs/src/DMC.md

@@ -0,0 +1,49 @@
+```@meta
+CurrentModule = TwoBody
+```
+
+# Diffusion Monte Carlo
+
+Pure diffusion Monte Carlo propagates walkers in imaginary time. For
+``H=-D\nabla^2+V``, one step combines Gaussian diffusion with branching weights
+
+```math
+w = \exp[-\Delta t(V-E_\mathrm{ref})].
+```
+
+`DiffusionMonteCarlo` keeps the population fixed by systematic resampling and
+estimates the ground-state energy from the reference-energy history. A seeded
+random-number generator makes calculations reproducible.
+
+## Usage
+
+```@example dmc
+using TwoBody
+
+H = Hamiltonian(
+  Kinetic(hbar=1, m=1),
+  PowerLaw(coefficient=1 / 2, exponent=2),
+)
+method = DiffusionMonteCarlo(
+  n_steps=800,
+  equilibration=200,
+  n_walkers=1_000,
+  Δt=0.01,
+)
+result = solve(H, method)
+result.E
+```
+
+Finite time steps and walker populations introduce bias. Convergence should be
+checked by reducing ``\Delta t`` and increasing `n_walkers`. See
+[Reynolds et al. (1990)](https://doi.org/10.1063/1.4822960) and
+[Kosztin et al. (1996)](https://doi.org/10.1119/1.18168).
+The reported `standard_error` does not account for autocorrelation.
+
+## API reference
+
+```@docs; canonical=false
+DiffusionMonteCarlo
+ResultDiffusionMonteCarlo
+solve(hamiltonian::Hamiltonian, method::DiffusionMonteCarlo)
+```

src/DMC.jl

@@ -0,0 +1,180 @@
+export DiffusionMonteCarlo, ResultDiffusionMonteCarlo
+
+import LinearAlgebra
+import Random
+
+struct DiffusionMonteCarlo{T<:AbstractFloat}
+  n_steps::Int
+  equilibration::Int
+  n_walkers::Int
+  Δt::T
+  initial_scale::T
+  feedback::T
+
+  function DiffusionMonteCarlo(;
+    n_steps::Int=2_000,
+    equilibration::Int=500,
+    n_walkers::Int=1_000,
+    Δt::Real=0.01,
+    initial_scale::Real=1.0,
+    feedback::Real=0.1,
+  )
+    0 < n_steps || throw(ArgumentError("n_steps must be positive"))
+    0 ≤ equilibration < n_steps ||
+      throw(ArgumentError("equilibration must satisfy 0 ≤ equilibration < n_steps"))
+    0 < n_walkers || throw(ArgumentError("n_walkers must be positive"))
+    isfinite(Δt) && 0 < Δt || throw(ArgumentError("Δt must be positive and finite"))
+    isfinite(initial_scale) && 0 < initial_scale ||
+      throw(ArgumentError("initial_scale must be positive and finite"))
+    isfinite(feedback) && 0 < feedback ≤ 1 ||
+      throw(ArgumentError("feedback must satisfy 0 < feedback ≤ 1"))
+    type = promote_type(typeof(float(Δt)), typeof(float(initial_scale)), typeof(float(feedback)))
+    new{type}(n_steps, equilibration, n_walkers, Δt, initial_scale, feedback)
+  end
+end
+
+struct ResultDiffusionMonteCarlo
+  data::Any
+  ResultDiffusionMonteCarlo(; args...) = new(NamedTuple(Dict(args)))
+end
+
+Base.getproperty(result::ResultDiffusionMonteCarlo, symbol::Symbol) =
+  Base.getproperty(getfield(result, :data), symbol)
+
+Base.string(method::DiffusionMonteCarlo) =
+  "DiffusionMonteCarlo(" *
+  join(["$(symbol)=$(getproperty(method, symbol))" for symbol in fieldnames(typeof(method))], ", ") *
+  ")"
+
+function Base.string(result::ResultDiffusionMonteCarlo)
+  return "# method\n\n$(result.method)\n\n# energy\n\nE = $(result.E)\n"
+end
+
+Base.show(io::IO, method::DiffusionMonteCarlo) = print(io, Base.string(method))
+Base.show(io::IO, result::ResultDiffusionMonteCarlo) = print(io, Base.string(result))
+
+_dmc_diffusion(term::Kinetic) = term.hbar^2 / (2 * term.m)
+_dmc_diffusion(term::Laplacian) = -term.coefficient
+_dmc_diffusion(::RestEnergy) = 0.0
+
+function _dmc_diffusion(term::KineticTerm)
+  throw(ArgumentError("$(typeof(term)) is not supported by DiffusionMonteCarlo"))
+end
+
+_dmc_potential(term::RestEnergy, radius) = term.m * term.c^2
+_dmc_potential(term::PotentialTerm, radius) = V(term, radius)
+
+function _dmc_potential(term::Union{Delta,Tabulated}, radius)
+  throw(ArgumentError("$(typeof(term)) is not supported by DiffusionMonteCarlo"))
+end
+
+_dmc_potential(::KineticTerm, radius) = 0.0
+
+function _dmc_potential(hamiltonian::Hamiltonian, walkers::AbstractMatrix)
+  potential = zeros(eltype(walkers), size(walkers, 2))
+  for term in hamiltonian.terms
+    for index in axes(walkers, 2)
+      potential[index] += _dmc_potential(term, LinearAlgebra.norm(@view walkers[:,index]))
+    end
+  end
+  all(isfinite, potential) ||
+    throw(ArgumentError("the potential energy must be finite at every walker"))
+  return potential
+end
+
+function _dmc_resample(weights, rng::Random.AbstractRNG)
+  total = sum(weights)
+  isfinite(total) && 0 < total || throw(ArgumentError("branching weights must have a finite positive sum"))
+  cumulative = cumsum(weights ./ total)
+  n_walkers = length(weights)
+  offset = rand(rng) / n_walkers
+  indices = Vector{Int}(undef, n_walkers)
+  source = 1
+  for target in 1:n_walkers
+    position = offset + (target - 1) / n_walkers
+    while cumulative[source] < position
+      source += 1
+    end
+    indices[target] = source
+  end
+  return indices
+end
+
+function solve(
+  hamiltonian::Hamiltonian,
+  method::DiffusionMonteCarlo;
+  rng::Random.AbstractRNG=Random.MersenneTwister(123),
+  initial=nothing,
+)
+  diffusion = 0.0
+  for term in hamiltonian.terms
+    term isa KineticTerm && (diffusion += _dmc_diffusion(term))
+  end
+  isfinite(diffusion) && 0 < diffusion ||
+    throw(ArgumentError("the Hamiltonian must have a positive nonrelativistic diffusion coefficient"))
+
+  walkers = if isnothing(initial)
+    method.initial_scale .* randn(rng, typeof(method.Δt), 3, method.n_walkers)
+  else
+    positions = Matrix{typeof(method.Δt)}(initial)
+    size(positions) == (3, method.n_walkers) ||
+      throw(DimensionMismatch("initial must have size (3, $(method.n_walkers))"))
+    all(isfinite, positions) || throw(ArgumentError("initial positions must be finite"))
+    positions
+  end
+
+  potential = _dmc_potential(hamiltonian, walkers)
+  reference_energy = sum(potential) / length(potential)
+  history = Vector{typeof(reference_energy)}(undef, method.n_steps)
+  σ = sqrt(2 * diffusion * method.Δt)
+
+  for step in 1:method.n_steps
+    proposed = walkers .+ σ .* randn(rng, typeof(method.Δt), size(walkers))
+    proposed_potential = _dmc_potential(hamiltonian, proposed)
+    branching_potential = (potential .+ proposed_potential) ./ 2
+    log_weights = -method.Δt .* (branching_potential .- reference_energy)
+    shift = maximum(log_weights)
+    weights = exp.(log_weights .- shift)
+    log_mean_weight = shift + log(sum(weights) / length(weights))
+    growth_energy = reference_energy - log_mean_weight / method.Δt
+    reference_energy += method.feedback * (growth_energy - reference_energy)
+
+    indices = _dmc_resample(weights, rng)
+    walkers = proposed[:,indices]
+    potential = proposed_potential[indices]
+    history[step] = reference_energy
+  end
+
+  retained = @view history[method.equilibration + 1:end]
+  energy = sum(retained) / length(retained)
+  variance = length(retained) == 1 ? zero(energy) :
+    sum((value - energy)^2 for value in retained) / (length(retained) - 1)
+  standard_error = sqrt(variance / length(retained))
+  return ResultDiffusionMonteCarlo(;
+    hamiltonian, method, E=energy, variance, standard_error,
+    reference_energies=history, walkers,
+  )
+end
+
+@doc raw"""
+`DiffusionMonteCarlo(; n_steps=2000, equilibration=500, n_walkers=1000, Δt=0.01, initial_scale=1.0, feedback=0.1)`
+
+Configure DMC sampling parameters.
+""" DiffusionMonteCarlo
+
+@doc raw"""
+`ResultDiffusionMonteCarlo`
+
+Result of a diffusion Monte Carlo calculation.
+""" ResultDiffusionMonteCarlo
+
+@doc raw"""
+`solve(hamiltonian, method::DiffusionMonteCarlo; rng=Random.MersenneTwister(123), initial=nothing)`
+
+Run the calculation and return a `ResultDiffusionMonteCarlo`.
+""" solve(
+  hamiltonian::Hamiltonian,
+  method::DiffusionMonteCarlo;
+  rng::Random.AbstractRNG=Random.MersenneTwister(123),
+  initial=nothing,
+)

src/TwoBody.jl

@@ -16,5 +16,6 @@ include("./FDM.jl")
 include("./QTT.jl")
 include("./VNN.jl")
 include("./VMC.jl")
+include("./DMC.jl")
 
 end

test/DMC.jl

@@ -0,0 +1,40 @@
+@testset "DMC.jl" begin
+  H = Hamiltonian(
+    Kinetic(hbar=1, m=1),
+    PowerLaw(coefficient=1 / 2, exponent=2),
+  )
+  method = DiffusionMonteCarlo(
+    n_steps=1_600,
+    equilibration=400,
+    n_walkers=1_000,
+    Δt=0.01,
+    initial_scale=1.0,
+  )
+  result = solve(H, method; rng=MersenneTwister(123))
+  repeated = solve(H, method; rng=MersenneTwister(123))
+
+  @test result isa ResultDiffusionMonteCarlo
+  @test result.E ≈ 1.5 atol=0.04
+  @test result.E == repeated.E
+  @test result.walkers == repeated.walkers
+  @test length(result.reference_energies) == method.n_steps
+  @test size(result.walkers) == (3, method.n_walkers)
+  @test isfinite(result.standard_error)
+  @test occursin("# energy", string(result))
+
+  hydrogen = Hamiltonian(Kinetic(hbar=1, m=1), Coulomb(coefficient=-1))
+  hydrogen_method = DiffusionMonteCarlo(
+    n_steps=1_500,
+    equilibration=500,
+    n_walkers=1_000,
+    Δt=0.01,
+  )
+  @test solve(hydrogen, hydrogen_method).E ≈ -0.5 atol=0.04
+
+  @test_throws ArgumentError DiffusionMonteCarlo(n_steps=0)
+  @test_throws ArgumentError DiffusionMonteCarlo(equilibration=2_000)
+  @test_throws ArgumentError DiffusionMonteCarlo(n_walkers=0)
+  @test_throws ArgumentError DiffusionMonteCarlo(Δt=0)
+  @test_throws ArgumentError solve(Hamiltonian(Coulomb(coefficient=-1)), method)
+  @test_throws DimensionMismatch solve(H, method; initial=zeros(2, method.n_walkers))
+end

test/runtests.jl

@@ -18,4 +18,5 @@ using Random
   include("QTT.jl")
   include("VNN.jl")
   include("VMC.jl")
+  include("DMC.jl")
 end