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.

Standard model

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

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

-0.46877911188327553

The returned normalized radial wavefunction is callable.

result.wavefunction(1.0)

0.22348529294839425

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

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

TwoBody.VariationalNeuralNetwork — Type

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

Options for optimizing a Lux neural network as a radial trial wavefunction. VNN is an abbreviation for VariationalNeuralNetwork.

Pass an existing FiniteDifferenceMethod as fdm, or use the finite-difference keywords directly. architecture defines the hidden-layer widths of the standard Lux model used by solve(hamiltonian, method). A custom Lux model can instead be supplied with solve(hamiltonian, model, method).

source

TwoBody.solve — Method

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

Build the standard Lux model specified by method.architecture and minimize its finite-difference Rayleigh quotient. See the three-argument overload to supply a custom Lux model.

source

TwoBody.solve — Method

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

Minimize 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}\]

with respect to the parameters of a Lux model. The model receives the complete radial grid as a batch and must return one real value per grid point. trial can impose an envelope or boundary condition on the raw output. Pass returned parameters, states, and optionally optimizer_state to continue training.

The result contains the energy E, normalized grid values ψ, raw values raw_ψ, a callable wavefunction, Lux variables, optimizer state, energy history, n_iterations, and converged.

source