diff --git a/docs/make.jl b/docs/make.jl index bdf5daa..773ca92 100644 --- a/docs/make.jl +++ b/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", ] diff --git a/docs/src/DMC.md b/docs/src/DMC.md new file mode 100644 index 0000000..34de027 --- /dev/null +++ b/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) +``` diff --git a/src/DMC.jl b/src/DMC.jl new file mode 100644 index 0000000..db87dd3 --- /dev/null +++ b/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, +) diff --git a/src/TwoBody.jl b/src/TwoBody.jl index 72dbe16..37aae63 100644 --- a/src/TwoBody.jl +++ b/src/TwoBody.jl @@ -16,5 +16,6 @@ include("./FDM.jl") include("./QTT.jl") include("./VNN.jl") include("./VMC.jl") +include("./DMC.jl") end diff --git a/test/DMC.jl b/test/DMC.jl new file mode 100644 index 0000000..97b9461 --- /dev/null +++ b/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 diff --git a/test/runtests.jl b/test/runtests.jl index 712aa08..f27c453 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -18,4 +18,5 @@ using Random include("QTT.jl") include("VNN.jl") include("VMC.jl") + include("DMC.jl") end