Add variational neural network solver

Project.toml

@@ -8,20 +8,26 @@ ArnoldiMethod = "ec485272-7323-5ecc-a04f-4719b315124d"
 FiniteDifferenceMatrices = "a7a66f33-e7b8-47af-b618-f9b5bea05f3d"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
+Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 Subscripts = "2b7f82d5-8785-4f63-971e-f18ddbeb808e"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 ArnoldiMethod = "0.4.0"
 FiniteDifferenceMatrices = "0.1.0"
 ForwardDiff = "0.10, 1"
+Lux = "0.5, 1"
 Optim = "1.9.4"
+Optimisers = "0.2, 0.3, 0.4"
 QuadGK = "2.11"
 SpecialFunctions = "2.3.1"
 Subscripts = "0.1.3"
+Zygote = "0.6, 0.7"
 julia = "1.7"

README.md

@@ -18,13 +18,15 @@ flowchart TD
   A["Hamiltonian.jl"]
   C["Rayleigh-Ritz.jl"]
   F["FDM.jl"]
+  N["VNN.jl"]
   G["VMC.jl"]
   H["DB.jl"]
   Z["TwoBody.jl"]
   A --> H
-  A --> C & F & G
-  H --> C & F & G
-  C & F & G --> Z
+  A --> C & F & N & G
+  H --> C & F & N & G
+  F --> N
+  C & F & N & G --> Z
 ```
 
 ## Developer's Guide

docs/Project.toml

@@ -2,6 +2,8 @@
 Antique = "be6e5d0e-34a5-4c8f-af83-e1b5389203d8"
 CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
+Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 TwoBody = "a92d7657-722c-45a6-9d18-9da4c8a753b6"
 
 [compat]

docs/make.jl

@@ -22,6 +22,7 @@ makedocs(;
         "Rayleigh-Ritz Method" => "Rayleigh-Ritz.md",
         "Free Complement Method" => "Free-Complement.md",
         "Finite Difference Method" => "FDM.md",
+        "Variational Neural Network" => "VNN.md",
         "Variational Monte Carlo" => "VMC.md",
         "API reference" => "API.md",
     ],

docs/src/VNN.md

@@ -0,0 +1,103 @@
+```@meta
+CurrentModule = TwoBody
+```
+
+# Variational Neural Network
+
+`VariationalNeuralNetwork` (`VNN`) uses a neural network defined with
+[Lux.jl](https://lux.csail.mit.edu/) as a radial trial wavefunction. It minimizes
+the finite-difference Rayleigh quotient
+
+```math
+E[\psi_\theta] =
+\frac{\pmb{\psi}_\theta^\mathsf{T}\pmb{J}\pmb{H}\pmb{\psi}_\theta}
+     {\pmb{\psi}_\theta^\mathsf{T}\pmb{J}\pmb{\psi}_\theta},
+```
+
+where the grid, Hamiltonian matrix ``\pmb{H}``, and radial Jacobian ``\pmb{J}``
+are provided by `FiniteDifferenceMethod`.
+
+## Standard model
+
+The two-argument `solve` method constructs a Lux network from `architecture`.
+
+```@example vnn
+using TwoBody
+
+H = Hamiltonian(
+  Kinetic(hbar=1, m=1),
+  Coulomb(coefficient=-1),
+)
+
+method = VNN(
+  Δr=0.2,
+  rₘₐₓ=4.0,
+  architecture=[4],
+  maxiters=100,
+  every=25,
+  abstol=0,
+)
+
+result = solve(H, method; info=1)
+result.E
+```
+
+The returned normalized radial wavefunction is callable.
+
+```@example vnn
+result.wavefunction(1.0)
+```
+
+## Custom Lux model
+
+An arbitrary Lux model can be passed explicitly. The model receives radii as a
+`1 × number_of_grid_points` batch and must return one real value per radius.
+
+```julia
+using Lux
+using Optimisers
+using Random
+
+model = Lux.Chain(
+  Lux.Dense(1 => 8, tanh),
+  Lux.Dense(8 => 1),
+)
+
+method = VNN(
+  fdm=FiniteDifferenceMethod(Δr=0.1, rₘₐₓ=20.0),
+  optimizer=Optimisers.Adam(0.01),
+  maxiters=2_000,
+)
+
+result = solve(
+  H,
+  model,
+  method;
+  rng=Random.MersenneTwister(123),
+  trial=(r, value) -> exp(-r) * value,
+  info=1,
+)
+```
+
+`trial` can impose an envelope or boundary condition. To continue training,
+pass `result.parameters`, `result.states`, and optionally
+`result.optimizer_state` to another call. The result also contains normalized
+grid values `ψ`, raw values `raw_ψ`, `history`, `n_iterations`, and `converged`.
+
+## API reference
+
+```@docs; canonical=false
+TwoBody.VariationalNeuralNetwork
+TwoBody.solve(hamiltonian::Hamiltonian, method::VariationalNeuralNetwork)
+TwoBody.solve(
+  hamiltonian::Hamiltonian,
+  model,
+  method::VariationalNeuralNetwork;
+  rng::Random.AbstractRNG,
+  parameters,
+  states,
+  optimizer_state,
+  trial,
+  info::Int,
+)
+```

src/FDM.jl

@@ -12,7 +12,7 @@ import Printf
 struct FiniteDifferenceMethod
   Δr::Real
   rₘₐₓ::Real
-  R::StepRangeLen
+  R::AbstractRange
   l::Int
   direction::Symbol
   solver::Symbol
@@ -21,6 +21,23 @@ struct FiniteDifferenceMethod
   end
 end
 
+_jacobian(method::FiniteDifferenceMethod) = SparseArrays.spdiagm(method.R .^ 2)
+
+function _rayleigh_quotient(ψ::AbstractVector, H::AbstractMatrix, J::AbstractMatrix)
+  denominator = LinearAlgebra.dot(ψ, J * ψ)
+  return LinearAlgebra.dot(ψ, J * (H * ψ)) / denominator
+end
+
+function _normalization(ψ::AbstractVector, method::FiniteDifferenceMethod, J)
+  norm² = 4 * π * method.Δr * real(LinearAlgebra.dot(ψ, J * ψ))
+  isfinite(norm²) && 0 < norm² ||
+    throw(ArgumentError("the wavefunction norm must be positive and finite"))
+  return inv(sqrt(norm²))
+end
+
+_normalize_wavefunction(ψ::AbstractVector, method::FiniteDifferenceMethod, J) =
+  _normalization(ψ, method, J) * ψ
+
 Base.string(method::FiniteDifferenceMethod) = "FiniteDifferenceMethod(" * join(["$(symbol)=$(getproperty(method,symbol))" for symbol in fieldnames(typeof(method))], ", ") * ")"
 Base.show(io::IO, method::FiniteDifferenceMethod) = print(io, Base.string(method))
 
@@ -55,7 +72,7 @@ function solve(hamiltonian::Hamiltonian, method::FiniteDifferenceMethod; perturb
   H = matrix(hamiltonian, method)
 
   # Jacobian
-  J = SparseArrays.spdiagm(method.R .^ 2)
+  J = _jacobian(method)
 
   # Eigenvalues
   if method.solver == :LinearAlgebra
@@ -138,13 +155,13 @@ function solve(hamiltonian::Hamiltonian, wavefunction::Function, method::FiniteD
   H = matrix(hamiltonian, method)
 
   # Jacobian
-  J = SparseArrays.spdiagm(method.R .^ 2)
+  J = _jacobian(method)
 
   # Wave Function
   ψ = wavefunction.(method.R)
 
   # Energy
-  E = (ψ' * J * H * ψ) / (ψ' * J * ψ)
+  E = _rayleigh_quotient(ψ, H, J)
 
   # Return
   if 0 ≤ info
@@ -155,7 +172,7 @@ function solve(hamiltonian::Hamiltonian, wavefunction::Function, method::FiniteD
       H = H,
       J = J,
       E = E,
-      ψ = ψ / sqrt(4 * π * method.Δr * ψ' * J * ψ),
+      ψ = _normalize_wavefunction(ψ, method, J),
     )
   else
     return (
@@ -174,7 +191,7 @@ end
 | :-------- | :------ | :---------- |
 | `Δr::Real`          | `0.1`            | Radial grid spacing. A uniform grid spacing is used, ``r_{i+1} = r_{i} + \Delta r``. |
 | `rₘₐₓ::Real`        | `50.0`           | The maximum value of the radial grid. This value is not directly used in the calculation, but it is used to determine the `R`. |
-| `R::StepRangeLen`   | `Δr:Δr:rₘₐₓ`     | Radial grid. The origin must be excluded from the grid to avoid divergence of the Coulomb potential and the centrifugal potential at the origin. |
+| `R::AbstractRange`  | `Δr:Δr:rₘₐₓ`     | Radial grid. The origin must be excluded from the grid to avoid divergence of the Coulomb potential and the centrifugal potential at the origin. |
 | `l::Int`            | `0`              | Angular momentum quantum number. This is a positive integer, ``0 \leq l``. |
 | `direction::Symbol` | `:c`             | The direction of the finite difference, `:c` for central, :f for forward, `:b` for backward. |
 | `solver::Symbol`    | `:LinearAlgebra` | The solver for eigenvalue problem, `:LinearAlgebra` or `:ArnoldiMethod`. |

src/TwoBody.jl

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