From e69684e12668210403f7e3bbe6db11f51a376330 Mon Sep 17 00:00:00 2001 From: Shuhei Ohno Date: Tue, 11 Aug 2026 22:56:45 +0900 Subject: [PATCH] Add variational neural network solver --- Project.toml | 6 + README.md | 8 +- docs/Project.toml | 2 + docs/make.jl | 1 + docs/src/VNN.md | 103 +++++++++++++++ src/FDM.jl | 29 ++++- src/TwoBody.jl | 1 + src/VNN.jl | 319 ++++++++++++++++++++++++++++++++++++++++++++++ test/Project.toml | 2 + test/VNN.jl | 130 +++++++++++++++++++ test/runtests.jl | 1 + 11 files changed, 593 insertions(+), 9 deletions(-) create mode 100644 docs/src/VNN.md create mode 100644 src/VNN.jl create mode 100644 test/VNN.jl diff --git a/Project.toml b/Project.toml index 74c6852..27f6f28 100644 --- a/Project.toml +++ b/Project.toml @@ -8,18 +8,24 @@ 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" 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" SpecialFunctions = "2.3.1" Subscripts = "0.1.3" +Zygote = "0.6, 0.7" julia = "1.7" diff --git a/README.md b/README.md index a52e217..89bd1fb 100644 --- a/README.md +++ b/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 diff --git a/docs/Project.toml b/docs/Project.toml index e5fe655..3b2bed0 100644 --- a/docs/Project.toml +++ b/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] diff --git a/docs/make.jl b/docs/make.jl index 47bffd1..34329ce 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 Neural Network" => "VNN.md", "Variational Monte Carlo" => "VMC.md", "API reference" => "API.md", ], diff --git a/docs/src/VNN.md b/docs/src/VNN.md new file mode 100644 index 0000000..a7990f7 --- /dev/null +++ b/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, +) +``` diff --git a/src/FDM.jl b/src/FDM.jl index e8b8cec..6fe77c8 100644 --- a/src/FDM.jl +++ b/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`. | diff --git a/src/TwoBody.jl b/src/TwoBody.jl index 9a74c74..8e2607e 100644 --- a/src/TwoBody.jl +++ b/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 diff --git a/src/VNN.jl b/src/VNN.jl new file mode 100644 index 0000000..9cc6ea3 --- /dev/null +++ b/src/VNN.jl @@ -0,0 +1,319 @@ +export VariationalNeuralNetwork, VNN + +import Lux +import Optimisers +import Printf +import Random +import Zygote + +_softplus(x) = max(x, zero(x)) + log1p(exp(-abs(x))) + +struct VariationalNeuralNetwork{F<:FiniteDifferenceMethod,A,I,O,T<:AbstractFloat} + fdm::F + architecture::Vector{Int} + activation::A + init::I + optimizer::O + maxiters::Int + abstol::T + patience::Int + every::Int + + function VariationalNeuralNetwork(; + fdm::Union{Nothing,FiniteDifferenceMethod}=nothing, + Δr::Union{Nothing,Real}=nothing, + rₘₐₓ::Union{Nothing,Real}=nothing, + R::Union{Nothing,AbstractRange}=nothing, + l::Union{Nothing,Int}=nothing, + direction::Union{Nothing,Symbol}=nothing, + solver::Union{Nothing,Symbol}=nothing, + architecture::AbstractVector{<:Integer}=[2], + activation=_softplus, + init=Lux.glorot_normal, + optimizer=Optimisers.Adam(0.01), + maxiters::Int=1_000, + abstol::Real=1e-8, + patience::Int=10, + every::Int=100, + ) + if isnothing(fdm) + spacing = something(Δr, 0.1) + maximum_radius = something(rₘₐₓ, 50.0) + grid = isnothing(R) ? (spacing:spacing:maximum_radius) : R + fdm = FiniteDifferenceMethod( + Δr=spacing, + rₘₐₓ=maximum_radius, + R=grid, + l=something(l, 0), + direction=something(direction, :c), + solver=something(solver, :LinearAlgebra), + ) + elseif any(value -> !isnothing(value), (Δr, rₘₐₓ, R, l, direction, solver)) + throw(ArgumentError("pass either fdm or finite-difference keywords, not both")) + end + + _validate_vnn_fdm(fdm) + isempty(architecture) && throw(ArgumentError("architecture must not be empty")) + all(0 .< architecture) || + throw(ArgumentError("all hidden-layer widths must be positive")) + 0 <= maxiters || throw(ArgumentError("maxiters must be nonnegative")) + isfinite(abstol) && 0 <= abstol || + throw(ArgumentError("abstol must be nonnegative and finite")) + 0 < patience || throw(ArgumentError("patience must be positive")) + 0 < every || throw(ArgumentError("every must be positive")) + + tolerance = float(abstol) + new{ + typeof(fdm), + typeof(activation), + typeof(init), + typeof(optimizer), + typeof(tolerance), + }( + fdm, + Int.(architecture), + activation, + init, + optimizer, + maxiters, + tolerance, + patience, + every, + ) + end +end + +const VNN = VariationalNeuralNetwork + +function _validate_vnn_fdm(fdm::FiniteDifferenceMethod) + isfinite(fdm.Δr) && 0 < fdm.Δr || + throw(ArgumentError("fdm.Δr must be positive and finite")) + isfinite(fdm.rₘₐₓ) && 0 < fdm.rₘₐₓ || + throw(ArgumentError("fdm.rₘₐₓ must be positive and finite")) + isempty(fdm.R) && throw(ArgumentError("fdm.R must not be empty")) + all(isfinite, fdm.R) || throw(ArgumentError("fdm.R must contain only finite values")) + all(0 .< fdm.R) || throw(ArgumentError("fdm.R must contain only positive values")) + 0 <= fdm.l || throw(ArgumentError("fdm.l must be nonnegative")) + fdm.direction in (:c, :f, :b) || + throw(ArgumentError("fdm.direction must be :c, :f, or :b")) + return nothing +end + +Base.string(method::VariationalNeuralNetwork) = + "VariationalNeuralNetwork(" * + join(["$(symbol)=$(getproperty(method, symbol))" for symbol in fieldnames(typeof(method))], ", ") * + ")" +Base.show(io::IO, method::VariationalNeuralNetwork) = print(io, Base.string(method)) + +function _neural_network(method::VariationalNeuralNetwork) + dimensions = [1; method.architecture; 1] + layers = [ + Lux.Dense( + dimensions[index], + dimensions[index + 1], + method.activation; + init_weight=method.init, + ) for index in 1:(length(dimensions) - 1) + ] + return Lux.Chain(layers...) +end + +_parameter_eltype(x::AbstractArray) = eltype(x) +_parameter_eltype(x::Number) = typeof(x) + +function _parameter_eltype(x::Union{NamedTuple,Tuple}) + for value in values(x) + type = _parameter_eltype(value) + isnothing(type) || return type + end + return nothing +end + +_parameter_eltype(x) = nothing + +function _network_input(parameters, R) + type = _parameter_eltype(parameters) + isnothing(type) && throw(ArgumentError("the Lux model must contain at least one parameter")) + type <: Real || throw(ArgumentError("the Lux model parameters must be real-valued")) + return reshape(type.(R), 1, :) +end + +function _trial_values(model, parameters, states, input, trial) + output, updated_states = Lux.apply(model, input, parameters, states) + length(output) == length(input) || throw(DimensionMismatch( + "the Lux model must return one wavefunction value for each radial-grid point", + )) + values = trial.(vec(input), vec(output)) + eltype(values) <: Real || + throw(ArgumentError("the trial wavefunction must be real-valued")) + return values, updated_states +end + +function _validate_trial(values, energy) + all(isfinite, values) || + throw(ArgumentError("the trial wavefunction must contain only finite values")) + any(value -> !iszero(value), values) || + throw(ArgumentError("the trial wavefunction must not be identically zero")) + isfinite(energy) || throw(ArgumentError("the variational energy must be finite")) + return nothing +end + +function solve( + hamiltonian::Hamiltonian, + method::VariationalNeuralNetwork; + kwargs..., +) + return solve(hamiltonian, _neural_network(method), method; kwargs...) +end + +function solve( + hamiltonian::Hamiltonian, + model, + method::VariationalNeuralNetwork; + rng::Random.AbstractRNG=Random.MersenneTwister(123), + parameters=nothing, + states=nothing, + optimizer_state=nothing, + trial=(r, value) -> value, + info::Int=0, +) + !isnothing(optimizer_state) && isnothing(parameters) && + throw(ArgumentError("parameters must be supplied with optimizer_state")) + + if isnothing(parameters) || isnothing(states) + initialized_parameters, initialized_states = Lux.setup(rng, model) + parameters = isnothing(parameters) ? initialized_parameters : parameters + states = isnothing(states) ? initialized_states : states + end + + fdm = method.fdm + H = matrix(hamiltonian, fdm) + J = _jacobian(fdm) + input = _network_input(parameters, fdm.R) + + values, states = _trial_values(model, parameters, states, input, trial) + energy = _rayleigh_quotient(values, H, J) + _validate_trial(values, energy) + + history = [energy] + optimizer_state = isnothing(optimizer_state) ? + Optimisers.setup(method.optimizer, parameters) : optimizer_state + converged = false + stable_steps = 0 + n_iterations = 0 + + if 0 < info + println("\n# method\n") + println(method) + println("\n# optimization\n") + Printf.@printf("%9s\t%14s\n", "iteration", "energy") + Printf.@printf("%9d\t%+.12e\n", 0, energy) + end + + for iteration in 1:method.maxiters + gradients = first(Zygote.gradient(parameters) do candidate_parameters + candidate_values, _ = + _trial_values(model, candidate_parameters, states, input, trial) + _rayleigh_quotient(candidate_values, H, J) + end) + optimizer_state, parameters = + Optimisers.update(optimizer_state, parameters, gradients) + + previous_energy = energy + values, states = _trial_values(model, parameters, states, input, trial) + energy = _rayleigh_quotient(values, H, J) + _validate_trial(values, energy) + push!(history, energy) + n_iterations = iteration + + if abs(energy - previous_energy) <= method.abstol + stable_steps += 1 + converged = method.patience <= stable_steps + else + stable_steps = 0 + end + + if 0 < info && (iteration % method.every == 0 || converged || iteration == method.maxiters) + Printf.@printf("%9d\t%+.12e\n", iteration, energy) + end + converged && break + end + + normalization = _normalization(values, fdm, J) + normalized_values = normalization * values + parameter_type = _parameter_eltype(parameters) + wavefunction = function (r::Real) + scalar_input = reshape(parameter_type[r], 1, 1) + output, _ = Lux.apply(model, scalar_input, parameters, states) + return normalization * trial(r, only(output)) + end + + return ( + hamiltonian=hamiltonian, + method=method, + model=model, + parameters=parameters, + states=states, + optimizer_state=optimizer_state, + H=H, + J=J, + E=energy, + ψ=normalized_values, + raw_ψ=values, + wavefunction=wavefunction, + history=history, + n_iterations=n_iterations, + converged=converged, + ) +end + +@doc raw""" +`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)`. +""" VariationalNeuralNetwork + +@doc raw""" +`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. +""" solve(hamiltonian::Hamiltonian, method::VariationalNeuralNetwork; kwargs...) + +@doc raw""" +`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 + +```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} +``` + +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`. +""" solve( + hamiltonian::Hamiltonian, + model, + method::VariationalNeuralNetwork; + rng::Random.AbstractRNG=Random.MersenneTwister(123), + parameters=nothing, + states=nothing, + optimizer_state=nothing, + trial=(r, value) -> value, + info::Int=0, +) diff --git a/test/Project.toml b/test/Project.toml index b1befd1..9996b77 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,6 +1,8 @@ [deps] Antique = "be6e5d0e-34a5-4c8f-af83-e1b5389203d8" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +Lux = "b2108857-7c20-44ae-9111-449ecde12c47" +Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" diff --git a/test/VNN.jl b/test/VNN.jl new file mode 100644 index 0000000..19c5a9c --- /dev/null +++ b/test/VNN.jl @@ -0,0 +1,130 @@ +@testset "VNN.jl" begin + import Lux + import Optimisers + + @testset "method validation" begin + @test VNN === VariationalNeuralNetwork + @test_throws ArgumentError VNN(Δr=0) + @test_throws ArgumentError VNN(rₘₐₓ=0) + @test_throws ArgumentError VNN(l=-1) + @test_throws ArgumentError VNN(direction=:invalid) + @test_throws ArgumentError VNN(architecture=Int[]) + @test_throws ArgumentError VNN(architecture=[2, 0]) + @test_throws ArgumentError VNN(maxiters=-1) + @test_throws ArgumentError VNN(abstol=-1) + @test_throws ArgumentError VNN(abstol=Inf) + @test_throws ArgumentError VNN(patience=0) + @test_throws ArgumentError VNN(every=0) + @test_throws ArgumentError VNN(fdm=FiniteDifferenceMethod(), Δr=0.2) + end + + @testset "standard model" begin + hamiltonian = Hamiltonian( + Kinetic(hbar=1, m=1), + PowerLaw(coefficient=1 / 2, exponent=2), + ) + method = VNN( + Δr=0.2, + rₘₐₓ=2.0, + architecture=[2], + maxiters=2, + abstol=0, + every=1, + ) + result = solve(hamiltonian, method) + repeated = solve(hamiltonian, method) + + @test isfinite(result.E) + @test result.history == repeated.history + @test result.E == last(result.history) + @test result.n_iterations == 2 + @test length(result.history) == result.n_iterations + 1 + @test size(result.ψ) == size(method.fdm.R) + @test 4π * method.fdm.Δr * (result.ψ' * result.J * result.ψ) ≈ 1 atol=1e-6 + @test isfinite(result.wavefunction(0.5)) + end + + @testset "custom hydrogen trial wavefunction" begin + hamiltonian = Hamiltonian( + Kinetic(hbar=1, m=1), + Coulomb(coefficient=-1), + ) + fdm = FiniteDifferenceMethod(Δr=0.1, rₘₐₓ=10.0) + initial_weight = (rng, output_dimension, input_dimension) -> + fill(0.6f0, output_dimension, input_dimension) + model = Lux.Dense( + 1 => 1, + value -> exp(-abs(value)); + use_bias=false, + init_weight=initial_weight, + ) + + evaluation_method = VNN(fdm=fdm, maxiters=0) + evaluation = solve(hamiltonian, model, evaluation_method) + reference = solve(hamiltonian, r -> exp(-0.6f0 * r), fdm, 0, 1) + + @test evaluation.E ≈ reference.E rtol=1e-6 + @test evaluation.n_iterations == 0 + @test !evaluation.converged + @test length(evaluation.history) == 1 + @test 4π * fdm.Δr * sum(fdm.R .^ 2 .* abs2.(evaluation.ψ)) ≈ 1 atol=1e-12 + @test evaluation.wavefunction(1.0) ≈ evaluation.ψ[10] rtol=1e-6 + + training_method = VNN( + fdm=fdm, + optimizer=Optimisers.Adam(0.03), + maxiters=100, + abstol=1e-7, + patience=5, + every=25, + ) + result = solve(hamiltonian, model, training_method) + + @test result.E < first(result.history) + @test result.E < -0.49 + @test result.parameters.weight[1] ≈ 1 atol=0.02 + @test length(result.history) == result.n_iterations + 1 + @test result.n_iterations <= training_method.maxiters + + resumed = solve( + hamiltonian, + model, + evaluation_method; + parameters=result.parameters, + states=result.states, + ) + @test resumed.E ≈ result.E rtol=1e-10 + + continued = solve( + hamiltonian, + model, + VNN(fdm=fdm, optimizer=training_method.optimizer, maxiters=1, abstol=0); + parameters=result.parameters, + states=result.states, + optimizer_state=result.optimizer_state, + ) + @test continued.n_iterations == 1 + end + + @testset "invalid model output" begin + zero_initializer = (rng, output_dimension, input_dimension) -> + zeros(Float32, output_dimension, input_dimension) + zero_model = Lux.Dense( + 1 => 1; + use_bias=false, + init_weight=zero_initializer, + ) + wrong_size_model = Lux.Dense(1 => 2) + method = VNN(Δr=0.5, rₘₐₓ=2.0, maxiters=0) + hamiltonian = Hamiltonian(Kinetic()) + + @test_throws ArgumentError solve(hamiltonian, zero_model, method) + @test_throws DimensionMismatch solve(hamiltonian, wrong_size_model, method) + @test_throws ArgumentError solve( + hamiltonian, + zero_model, + method; + optimizer_state=NamedTuple(), + ) + end +end diff --git a/test/runtests.jl b/test/runtests.jl index cc66f2c..4de449a 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -12,5 +12,6 @@ using ForwardDiff include("Basis.jl") include("Rayleigh-Ritz.jl") include("FDM.jl") + include("VNN.jl") include("VMC.jl") end