diff --git a/Project.toml b/Project.toml index 25ffb95..74c6852 100644 --- a/Project.toml +++ b/Project.toml @@ -6,9 +6,11 @@ version = "0.0.10" [deps] ArnoldiMethod = "ec485272-7323-5ecc-a04f-4719b315124d" FiniteDifferenceMatrices = "a7a66f33-e7b8-47af-b618-f9b5bea05f3d" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" Optim = "429524aa-4258-5aef-a3af-852621145aeb" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" Subscripts = "2b7f82d5-8785-4f63-971e-f18ddbeb808e" @@ -16,6 +18,7 @@ Subscripts = "2b7f82d5-8785-4f63-971e-f18ddbeb808e" [compat] ArnoldiMethod = "0.4.0" FiniteDifferenceMatrices = "0.1.0" +ForwardDiff = "0.10, 1" Optim = "1.9.4" SpecialFunctions = "2.3.1" Subscripts = "0.1.3" diff --git a/docs/Project.toml b/docs/Project.toml index 106e3c9..e5fe655 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -3,3 +3,6 @@ Antique = "be6e5d0e-34a5-4c8f-af83-e1b5389203d8" CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" TwoBody = "a92d7657-722c-45a6-9d18-9da4c8a753b6" + +[compat] +Antique = "0.15" diff --git a/docs/make.jl b/docs/make.jl index f830c53..2b77ee0 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -21,6 +21,7 @@ makedocs(; "Database" => "DB.md", "Rayleigh-Ritz Method" => "Rayleigh-Ritz.md", "Finite Difference Method" => "FDM.md", + "Variational Monte Carlo" => "VMC.md", "API reference" => "API.md", ], ) diff --git a/docs/src/VMC.md b/docs/src/VMC.md new file mode 100644 index 0000000..f07c22f --- /dev/null +++ b/docs/src/VMC.md @@ -0,0 +1,112 @@ +```@meta +CurrentModule = TwoBody +``` + +# Variational Monte Carlo + +VMC estimates ``\langle E\rangle`` by averaging the local energy over positions sampled from the normalized density ``P(\mathbf{r})``: + +```math +\begin{aligned} +\langle E \rangle +&= \frac{ + \displaystyle \int \mathrm{d}\mathbf{r}\, + \psi^*(\mathbf{r})\hat{H}\psi(\mathbf{r}) +}{ + \displaystyle \int \mathrm{d}\mathbf{r}\, + |\psi(\mathbf{r})|^2 +} \\ +&= \frac{1}{Z}\int \mathrm{d}\mathbf{r}\, +\psi^*(\mathbf{r})\hat{H}\psi(\mathbf{r}) +&& \qquad\cdots\qquad {\textstyle +Z = \int \mathrm{d}\mathbf{r}\,|\psi(\mathbf{r})|^2} \\ +&= \int \mathrm{d}\mathbf{r}\, +\frac{1}{Z}|\psi(\mathbf{r})|^2 +\cdot +\frac{\hat{H}\psi(\mathbf{r})}{\psi(\mathbf{r})} +&& \qquad\cdots\qquad {\textstyle +|\psi(\mathbf{r})|^2 = \psi^*(\mathbf{r})\psi(\mathbf{r})} \\ +&= \int \mathrm{d}\mathbf{r}\,P(\mathbf{r})\cdot E_\mathrm{loc}(\mathbf{r}) +&& \qquad\cdots\qquad {\textstyle +P(\mathbf{r}) = \frac{1}{Z}|\psi(\mathbf{r})|^2,\quad +E_\mathrm{loc}(\mathbf{r}) = +\frac{\hat{H}\psi(\mathbf{r})}{\psi(\mathbf{r})}} +\\ +&\approx \frac{1}{N}\sum_{i=1}^{N}E_\mathrm{loc}(\mathbf{r}_i) +&& \qquad\cdots\qquad {\textstyle \mathbf{r}_i \sim P.} +\end{aligned} +``` + +This local-energy formulation is reviewed by Foulkes *et al.* +[1]. Samples are generated with a symmetric random-walk proposal and the +Metropolis acceptance rule [2]. The Laplacian in a non-relativistic kinetic term +is evaluated by forward-mode automatic differentiation with ForwardDiff.jl [4]. + +## Usage + +The following example uses the hydrogen trial wavefunction ``\psi(r)=\exp(-0.8r)`` and sampling conditions from Thijssen [5, Table 12.1]: 300 walkers, 12,000 attempted displacements per walker, and the first 2,000 states of each walker discarded for equilibration (`burn_in=2_000`; `n_steps=10_000` counts the retained samples). A nonzero initial position avoids starting exactly at the Coulomb singularity. + +```@example vmc-hydrogen +using TwoBody + +# Hamiltonian +H = Hamiltonian( + NonRelativisticKinetic(ℏ=1, m=1), + CoulombPotential(coefficient=-1), +) + +# Trial wave function +α = 0.8 +ψ(r) = exp(-α * sqrt(sum(abs2, r))) + +# VMC options +method = VariationalMonteCarlo( + n_walkers=300, + n_steps=10_000, + burn_in=2_000, + thinning=1, + δ=2.0, + r₀=[1.0, 0.0, 0.0], +) + +# Solve +result = solve(H, ψ, method) + +# Display +println("This work: $(result.E)") +println("Reference: -0.4813(6)") +println("Exact : -0.4800") +``` + +For ``\psi(r)=\exp(-\alpha r)``, the analytical expectation value in atomic units is ``\alpha^2/2-\alpha=-0.48`` at ``\alpha=0.8``; Thijssen reports ``-0.4813(6)``. A finite-sample estimate fluctuates around the analytical expectation value; more samples reduce statistical noise but not the trial wavefunction's variational bias. + +`solve` also returns sampling diagnostics and retained data; see the [API reference](#API-reference) for details. + +Because successive Markov-chain samples are correlated, the naive `standard_error` reported by `solve` typically underestimates the true uncertainty; rigorous estimates require batching or autocorrelation analysis, such as the blocking method of Flyvbjerg and Petersen [3]. + +## Bibliography + +1. W. M. C. Foulkes, L. Mitas, R. J. Needs, and G. Rajagopal, + ["Quantum Monte Carlo simulations of solids"](https://doi.org/10.1103/RevModPhys.73.33), + *Reviews of Modern Physics* **73**, 33–83 (2001). +2. N. Metropolis, A. W. Rosenbluth, M. N. Rosenbluth, A. H. Teller, and E. Teller, + ["Equation of State Calculations by Fast Computing Machines"](https://doi.org/10.1063/1.1699114), + *The Journal of Chemical Physics* **21**, 1087–1092 (1953). +3. H. Flyvbjerg and H. G. Petersen, + ["Error estimates on averages of correlated data"](https://doi.org/10.1063/1.457480), + *The Journal of Chemical Physics* **91**, 461–466 (1989). +4. J. Revels, M. Lubin, and T. Papamarkou, + ["Forward-Mode Automatic Differentiation in Julia"](https://arxiv.org/abs/1607.07892), + arXiv:1607.07892 (2016). +5. J. M. Thijssen, + [*Computational Physics*, 2nd ed.](https://doi.org/10.1017/CBO9781139171397), + Cambridge University Press (2007); + 邦訳: 松田和典, 道廣嘉隆, 谷村吉隆, 高須昌子, 吉江友照 訳, 『計算物理学』, 丸善出版 (2012). + +## API reference + +```@docs; canonical=false +TwoBody.VariationalMonteCarlo +TwoBody.local_energy +TwoBody.solve(hamiltonian::Hamiltonian, wavefunction::Function, method::VariationalMonteCarlo) +``` diff --git a/src/TwoBody.jl b/src/TwoBody.jl index 37adc95..9a74c74 100644 --- a/src/TwoBody.jl +++ b/src/TwoBody.jl @@ -12,5 +12,6 @@ include("./Basis.jl") # Solvers include("./Rayleigh-Ritz.jl") include("./FDM.jl") +include("./VMC.jl") end diff --git a/src/VMC.jl b/src/VMC.jl new file mode 100644 index 0000000..1071baf --- /dev/null +++ b/src/VMC.jl @@ -0,0 +1,234 @@ +export VariationalMonteCarlo, local_energy + +import ForwardDiff +import LinearAlgebra +import Random + +struct VariationalMonteCarlo{T<:AbstractFloat,V<:AbstractVector{T}} + n_steps::Int + burn_in::Int + thinning::Int + n_walkers::Int + δ::T + r₀::V + + function VariationalMonteCarlo(; + n_steps::Int=10^5, + burn_in::Int=10^3, + thinning::Int=1, + n_walkers::Int=1, + δ::Real=0.5, + r₀::AbstractVector{<:Real}=[1.0, 0.0, 0.0], + ) + 0 < n_steps || throw(ArgumentError("n_steps must be positive")) + 0 ≤ burn_in || throw(ArgumentError("burn_in must be nonnegative")) + 0 < thinning || throw(ArgumentError("thinning must be positive")) + 0 < n_walkers || throw(ArgumentError("n_walkers must be positive")) + isfinite(δ) && 0 < δ || throw(ArgumentError("δ must be positive and finite")) + isempty(r₀) && throw(ArgumentError("r₀ must contain at least one coordinate")) + all(isfinite, r₀) || throw(ArgumentError("r₀ must contain only finite coordinates")) + + position = float.(collect(r₀)) + step_size = convert(eltype(position), δ) + new{eltype(position),typeof(position)}( + n_steps, + burn_in, + thinning, + n_walkers, + step_size, + position, + ) + end +end + +Base.string(method::VariationalMonteCarlo) = + "VariationalMonteCarlo(" * + join(["$(symbol)=$(getproperty(method, symbol))" for symbol in fieldnames(typeof(method))], ", ") * + ")" +Base.show(io::IO, method::VariationalMonteCarlo) = print(io, Base.string(method)) + +function _laplacian(wavefunction::Function, position::AbstractVector{<:Real}) + return LinearAlgebra.tr(ForwardDiff.hessian(wavefunction, position)) +end + +function _local_energy(term::Laplacian, wavefunction::Function, position, ψ) + return term.coefficient * _laplacian(wavefunction, position) / ψ +end + +function _local_energy(term::NonRelativisticKinetic, wavefunction::Function, position, ψ) + return -term.ℏ^2 / (2 * term.m) * _laplacian(wavefunction, position) / ψ +end + +_local_energy(term::RestEnergy, wavefunction::Function, position, ψ) = term.m * term.c^2 +_local_energy(term::PotentialTerm, wavefunction::Function, position, ψ) = + V(term, LinearAlgebra.norm(position)) + +function _local_energy(term::KineticTerm, wavefunction::Function, position, ψ) + throw(ArgumentError("$(typeof(term)) is not supported by VariationalMonteCarlo")) +end + +function local_energy( + hamiltonian::Hamiltonian, + wavefunction::Function, + position::AbstractVector{<:Real}, +) + ψ = wavefunction(position) + return sum(_local_energy(term, wavefunction, position, ψ) for term in hamiltonian.terms) +end + +function _probability_density(wavefunction::Function, position) + probability = abs2(wavefunction(position)) + probability isa Real || throw(ArgumentError("abs2(wavefunction(r)) must be real")) + return probability +end + +function _metropolis_samples( + wavefunction::Function, + method::VariationalMonteCarlo, + rng::Random.AbstractRNG, +) + total_samples = method.n_walkers * method.n_steps + samples = Matrix{eltype(method.r₀)}(undef, length(method.r₀), total_samples) + accepted = 0 + attempted = 0 + stored = 0 + n_transitions = method.burn_in + method.n_steps * method.thinning + half = eltype(method.r₀)(1 // 2) + + for _ in 1:method.n_walkers + position = copy(method.r₀) + probability = _probability_density(wavefunction, position) + isfinite(probability) && 0 < probability || + throw(ArgumentError("the probability density at r₀ must be positive and finite")) + + for step in 1:n_transitions + proposal = + position .+ method.δ .* (rand(rng, eltype(position), length(position)) .- half) + proposed_probability = _probability_density(wavefunction, proposal) + attempted += 1 + + if isfinite(proposed_probability) && 0 ≤ proposed_probability + acceptance_probability = min(one(probability), proposed_probability / probability) + if rand(rng) < acceptance_probability + position = proposal + probability = proposed_probability + accepted += 1 + end + end + + if method.burn_in < step && (step - method.burn_in) % method.thinning == 0 + stored += 1 + samples[:, stored] = position + end + end + end + + return samples, accepted, attempted +end + +function _isfinite_number(value::Number) + return isfinite(real(value)) && isfinite(imag(value)) +end + +function solve( + hamiltonian::Hamiltonian, + wavefunction::Function, + method::VariationalMonteCarlo; + rng::Random.AbstractRNG=Random.MersenneTwister(123), + info::Int=0, +) + samples, n_accepted, n_attempted = _metropolis_samples(wavefunction, method, rng) + acceptance_rate = n_accepted / n_attempted + energies = [local_energy(hamiltonian, wavefunction, samples[:, i]) for i in axes(samples, 2)] + finite_energies = filter(_isfinite_number, energies) + isempty(finite_energies) && + throw(ArgumentError("all local energies are non-finite; check the wavefunction and initial position")) + + energy = sum(finite_energies) / length(finite_energies) + variance = if length(finite_energies) == 1 + zero(abs2(first(finite_energies))) + else + sum(abs2(value - energy) for value in finite_energies) / (length(finite_energies) - 1) + end + standard_error = sqrt(variance / length(finite_energies)) + + result = ( + hamiltonian=hamiltonian, + method=method, + E=energy, + variance=variance, + standard_error=standard_error, + acceptance_rate=acceptance_rate, + n_accepted=n_accepted, + n_attempted=n_attempted, + n_burn_in_discarded=method.n_walkers * method.burn_in, + n_samples=length(finite_energies), + n_discarded=length(energies) - length(finite_energies), + local_energies=energies, + samples=samples, + ) + + if 0 < info + println("\n# method\n") + println(method) + println("\n# energy\n") + println("E = $(result.E) ± $(result.standard_error)") + println("\n# sampling\n") + println("acceptance rate = $(result.acceptance_rate)") + println("finite local energies = $(result.n_samples) / $(method.n_steps * method.n_walkers)") + end + + return result +end + +@doc raw""" +`VariationalMonteCarlo(n_steps=10^5, burn_in=10^3, thinning=1, n_walkers=1, δ=0.5, r₀=[1.0, 0.0, 0.0])` + +Options for variational Monte Carlo with a symmetric, uniform Metropolis proposal. +The sampler targets ``|\psi(\mathbf{r})|^2``. `n_steps` is the number of retained +samples per walker, `burn_in` is the number of initial transitions discarded from +each walker, `thinning` is the number of transitions between retained samples, +`n_walkers` is the number of Markov chains, `δ` is the proposal-box width, and +`r₀` is the initial position of every walker. Each walker performs +`burn_in + n_steps * thinning` transitions. +""" VariationalMonteCarlo + +@doc raw""" +`local_energy(hamiltonian, wavefunction, position)` + +Evaluate + +```math +E_\mathrm{loc}(\mathbf{r}) = +\frac{\hat{H}\psi(\mathbf{r})}{\psi(\mathbf{r})}. +``` + +The Laplacian of a real-valued trial wavefunction is evaluated with +`ForwardDiff.hessian`. Non-relativistic kinetic, Laplacian, rest-energy, and +potential terms with a defined `V` method are supported. +""" local_energy + +@doc raw""" +`solve(hamiltonian, wavefunction, method::VariationalMonteCarlo; rng=Random.MersenneTwister(123), info=0)` + +Estimate the variational energy by averaging the local energy over Metropolis +samples from ``|\psi|^2``. A seeded random-number generator can be supplied with +`rng`; by default, a new `MersenneTwister(123)` is used for a reproducible +calculation. Set `info` to a positive value to print a result summary. + +The returned named tuple echoes the input `hamiltonian` and `method` and +contains `E`, `variance`, the naive `standard_error`, +`acceptance_rate`, the sampling counts `n_accepted`, `n_attempted`, +`n_burn_in_discarded`, `n_samples`, and `n_discarded`, as well as +`local_energies` and `samples`. Data from multiple walkers are stored +consecutively, with positions in the columns of `samples`. Non-finite local +energies, which can occur at a measure-zero singularity such as the origin of a +Coulomb potential, are excluded from the energy statistics and counted in +`n_discarded`. +""" solve( + hamiltonian::Hamiltonian, + wavefunction::Function, + method::VariationalMonteCarlo; + rng::Random.AbstractRNG=Random.MersenneTwister(123), + info::Int=0, +) diff --git a/test/Project.toml b/test/Project.toml index 8c54d18..b1befd1 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -3,5 +3,9 @@ Antique = "be6e5d0e-34a5-4c8f-af83-e1b5389203d8" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" + +[compat] +Antique = "0.15" diff --git a/test/VMC.jl b/test/VMC.jl new file mode 100644 index 0000000..9079d43 --- /dev/null +++ b/test/VMC.jl @@ -0,0 +1,117 @@ +@testset "VMC.jl" begin + using Random + + @testset "method validation" begin + @test_throws ArgumentError VariationalMonteCarlo(n_steps=0) + @test_throws ArgumentError VariationalMonteCarlo(burn_in=-1) + @test_throws ArgumentError VariationalMonteCarlo(thinning=0) + @test_throws ArgumentError VariationalMonteCarlo(n_walkers=0) + @test_throws ArgumentError VariationalMonteCarlo(δ=0) + @test_throws ArgumentError VariationalMonteCarlo(δ=Inf) + @test_throws ArgumentError VariationalMonteCarlo(r₀=Float64[]) + end + + @testset "local energy" begin + H = Hamiltonian( + NonRelativisticKinetic(ℏ=1, m=1), + PowerLawPotential(coefficient=1 / 2, exponent=2), + ) + ψ(r) = exp(-sum(abs2, r) / 2) + + @test local_energy(H, ψ, [0.2, -0.3, 0.4]) ≈ 1.5 atol=1e-12 + @test_throws ArgumentError local_energy( + Hamiltonian(RelativisticKinetic()), + ψ, + [0.2, -0.3, 0.4], + ) + end + + @testset "harmonic oscillator" begin + H = Hamiltonian( + NonRelativisticKinetic(ℏ=1, m=1), + PowerLawPotential(coefficient=1 / 2, exponent=2), + ) + ψ(r) = exp(-sum(abs2, r) / 2) + method = VariationalMonteCarlo( + n_steps=200, + burn_in=20, + thinning=2, + δ=1.0, + r₀=[1.0, 0.0, 0.0], + ) + result = solve(H, ψ, method; rng=MersenneTwister(123), info=0) + + @test result.E ≈ 1.5 atol=1e-12 + @test result.n_samples == method.n_steps + @test result.n_discarded == 0 + @test size(result.samples) == (3, method.n_steps) + @test length(result.local_energies) == method.n_steps + @test 0 < result.acceptance_rate < 1 + @test result.n_attempted == method.burn_in + method.n_steps * method.thinning + mean_radius_squared = + sum(sum(abs2, result.samples[:, i]) for i in axes(result.samples, 2)) / method.n_steps + @test mean_radius_squared ≈ 1.5 atol=0.5 + + singular_H = Hamiltonian( + NonRelativisticKinetic(ℏ=1, m=1), + FunctionPotential(f=r -> r < 1 ? Inf : r^2 / 2), + ) + singular_result = solve( + singular_H, + ψ, + method; + rng=MersenneTwister(123), + info=0, + ) + @test 0 < singular_result.n_discarded < method.n_steps + @test isfinite(singular_result.E) + @test singular_result.n_samples + singular_result.n_discarded == method.n_steps + end + + @testset "multiple walkers" begin + H = Hamiltonian( + NonRelativisticKinetic(ℏ=1, m=1), + PowerLawPotential(coefficient=1 / 2, exponent=2), + ) + ψ(r) = exp(-sum(abs2, r) / 2) + method = VariationalMonteCarlo( + n_steps=50, + burn_in=10, + n_walkers=3, + δ=1.0, + r₀=[1.0, 0.0, 0.0], + ) + result = solve(H, ψ, method; rng=MersenneTwister(123), info=0) + total_samples = method.n_walkers * method.n_steps + transitions_per_walker = method.burn_in + method.n_steps * method.thinning + + @test result.E ≈ 1.5 atol=1e-12 + @test result.n_samples == total_samples + @test result.n_discarded == 0 + @test size(result.samples) == (3, total_samples) + @test length(result.local_energies) == total_samples + @test result.n_attempted == method.n_walkers * transitions_per_walker + @test result.n_accepted == round(Int, result.acceptance_rate * result.n_attempted) + @test result.n_burn_in_discarded == method.n_walkers * method.burn_in + @test 0 < result.acceptance_rate < 1 + + default_result = solve(H, ψ, method) + repeated_default_result = solve(H, ψ, method) + @test default_result.samples == repeated_default_result.samples + @test default_result.E == repeated_default_result.E + end + + @testset "Coulomb singularity" begin + H = Hamiltonian( + NonRelativisticKinetic(ℏ=1, m=1), + CoulombPotential(coefficient=-1), + ) + α = 0.2829 + ψ(r) = exp(-α * sum(abs2, r)) + method = VariationalMonteCarlo(n_steps=10_000, burn_in=2_000, δ=2.0) + result = solve(H, ψ, method; rng=MersenneTwister(123), info=0) + + @test isfinite(result.E) + @test result.E ≈ -0.42441317922678223 atol=0.03 + end +end diff --git a/test/runtests.jl b/test/runtests.jl index ad24e51..70262e0 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -11,4 +11,5 @@ using ForwardDiff include("Basis.jl") include("Rayleigh-Ritz.jl") include("FDM.jl") + include("VMC.jl") end