Add variational Monte Carlo

Project.toml

@@ -6,16 +6,19 @@ 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"
 
 [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"

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"

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",
     ],
 )

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)
+```

src/TwoBody.jl

@@ -12,5 +12,6 @@ include("./Basis.jl")
 # Solvers
 include("./Rayleigh-Ritz.jl")
 include("./FDM.jl")
+include("./VMC.jl")
 
 end

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,
+)

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"