Variational Neural Network

VariationalNeuralNetwork (VNN) uses a neural network defined with Lux.jl as a radial trial wavefunction. It minimizes the finite-difference Rayleigh quotient

\[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.

Install and load the optional dependencies to activate the VNN extension.

import Pkg; Pkg.add(["Lux", "Optimisers", "Zygote"])
using Lux, Optimisers, Zygote

Standard model

The two-argument solve method constructs a Lux network from architecture.

using TwoBody
using Lux
using Optimisers
using Zygote

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

-0.4687791161700152

The returned normalized radial wavefunction is callable.

result.wavefunction(1.0)

0.22348528635475506

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.

using Lux
using Optimisers
using Random
using Zygote

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

VariationalNeuralNetwork(; fdm=nothing, Δr=nothing, rₘₐₓ=nothing,
  R=nothing, l=nothing, direction=nothing, solver=nothing,
  architecture=[2], activation=softplus, init=nothing, optimizer=nothing,
  maxiters=1000, abstol=1e-8, patience=10, every=100)

solve(hamiltonian, method::VariationalNeuralNetwork; kwargs...)

solve(hamiltonian, model, method::VariationalNeuralNetwork;
      rng=Random.MersenneTwister(123), parameters=nothing, states=nothing,
      optimizer_state=nothing, trial=(r, value) -> value, info=0)