From b697f0e8bc0ceefc7f2188e5fccdd09261884fb6 Mon Sep 17 00:00:00 2001 From: MartinMikkelsen Date: Mon, 23 Feb 2026 11:53:19 +0100 Subject: [PATCH 1/8] init loss --- Examples/Variational.jl | 128 ++++++++++++++++ Project.toml | 4 + src/FewBodyECG.jl | 3 + src/hamiltonian.jl | 9 +- src/matrix_elements.jl | 6 +- src/variational.jl | 308 +++++++++++++++++++++++++++++++++++++++ test/runtests.jl | 1 + test/test_variational.jl | 218 +++++++++++++++++++++++++++ 8 files changed, 671 insertions(+), 6 deletions(-) create mode 100644 Examples/Variational.jl create mode 100644 src/variational.jl create mode 100644 test/test_variational.jl diff --git a/Examples/Variational.jl b/Examples/Variational.jl new file mode 100644 index 0000000..fd2c044 --- /dev/null +++ b/Examples/Variational.jl @@ -0,0 +1,128 @@ +# Variational ECG optimisation example +# +# Demonstrates solve_ECG_variational for two benchmark systems: +# +# 1. Hydrogen anion H⁻ (3-body: nucleus + 2 electrons) +# Exact ground-state energy: -0.527751016523 Ha +# +# 2. Muonic molecule tdμ (3-body: triton + deuteron + muon) +# Exact ground-state energy: -111.36444 Ha +# +# Two usage patterns are shown: +# +# (a) Fresh optimisation with loss_type = :energy (default) +# — starts from a QMC-generated initial basis. +# +# (b) Warm-start from solve_ECG result with loss_type = :trace +# — refines an already-good stochastic basis by minimising +# Tr(S⁻¹H). Requires the initial trace to be negative, +# which is guaranteed when the stochastic basis is close +# to the physical ground state. + +using FewBodyECG +using LinearAlgebra +using QuasiMonteCarlo +import FewBodyECG: default_scale, BasisSet, Rank0Gaussian + +# ============================================================ +# 1. Hydrogen anion H⁻ +# ============================================================ + +println("=" ^ 60) +println("Hydrogen anion H⁻") +println("=" ^ 60) + +masses_Hm = [1.0e15, 1.0, 1.0] # fixed nucleus + 2 electrons +Λ_Hm = Λ(masses_Hm) +_, U_Hm = _jacobi_transform(masses_Hm) + +w_pairs = [[1, -1, 0], [1, 0, -1], [0, 1, -1]] +w_raw_Hm = [U_Hm' * Float64.(w) for w in w_pairs] +coeffs_Hm = [-1.0, -1.0, +1.0] # e-nucleus (×2) and e-e repulsion + +ops_Hm = Operator[ + KineticOperator(Λ_Hm); + [CoulombOperator(c, w) for (c, w) in zip(coeffs_Hm, w_raw_Hm)]... +] + +E_exact_Hm = -0.527751016523 +n = 30 + +# --- (a) fresh optimisation with :energy loss --- +println("\n(a) Fresh start, loss_type = :energy") +sr_var = solve_ECG_variational(ops_Hm, n; scale = 1.0, max_iterations = 500, verbose = false) +ΔE_var = sr_var.ground_state - E_exact_Hm +println(" Variational E₀ = $(round(sr_var.ground_state, digits=8)) ΔE = $(round(ΔE_var, sigdigits=3))") + +# --- stochastic baseline --- +sr_stoch = solve_ECG(ops_Hm, n; scale = 1.0, verbose = false) +ΔE_stoch = sr_stoch.ground_state - E_exact_Hm +println(" Stochastic E₀ = $(round(sr_stoch.ground_state, digits=8)) ΔE = $(round(ΔE_stoch, sigdigits=3))") +println(" Exact E₀ = $E_exact_Hm") + +# --- (b) warm-start from stochastic result with :trace loss --- +println("\n(b) Warm-start from stochastic, loss_type = :trace") +basis0_Hm = BasisSet(Rank0Gaussian[sr_stoch.basis_functions...]) +sr_warm = solve_ECG_variational(ops_Hm, n; + initial_basis = basis0_Hm, + loss_type = :trace, + max_iterations = 500, + verbose = false, +) +ΔE_warm = sr_warm.ground_state - E_exact_Hm +println(" Warm-start E₀ = $(round(sr_warm.ground_state, digits=8)) ΔE = $(round(ΔE_warm, sigdigits=3))") +println(" Stochastic E₀ = $(round(sr_stoch.ground_state, digits=8)) ΔE = $(round(ΔE_stoch, sigdigits=3))") +println(" Exact E₀ = $E_exact_Hm") + +# --- downstream utilities work unchanged --- +r_grid, ρ = correlation_function(sr_var; rmin = 0.01, rmax = 15.0, npoints = 200) +println("\n Correlation function computed: $(length(r_grid)) points, max ρ at r = $(round(r_grid[argmax(ρ)], digits=3)) a.u.") + +# ============================================================ +# 2. Muonic molecule tdμ +# ============================================================ + +println() +println("=" ^ 60) +println("Muonic molecule tdμ (triton + deuteron + muon)") +println("=" ^ 60) + +masses_tdμ = [5496.918, 3670.481, 206.7686] # t, d, μ in electron masses +Λ_tdμ = Λ(masses_tdμ) +_, U_tdμ = _jacobi_transform(masses_tdμ) + +w_raw_tdμ = [U_tdμ' * Float64.(w) for w in w_pairs] +coeffs_tdμ = [+1.0, -1.0, -1.0] # t-d repulsion, t-μ and d-μ attraction + +ops_tdμ = Operator[ + KineticOperator(Λ_tdμ); + [CoulombOperator(c, w) for (c, w) in zip(coeffs_tdμ, w_raw_tdμ)]... +] + +E_exact_tdμ = -111.36444 +scale_tdμ = 0.03 # nuclear scale (much smaller than atomic) +n_tdμ = 25 + +println("\n(a) Fresh start, loss_type = :energy") +sr_var_tdμ = solve_ECG_variational(ops_tdμ, n_tdμ; + scale = scale_tdμ, max_iterations = 500, verbose = false +) +ΔE_var_tdμ = sr_var_tdμ.ground_state - E_exact_tdμ +println(" Variational E₀ = $(round(sr_var_tdμ.ground_state, digits=4)) ΔE = $(round(ΔE_var_tdμ, sigdigits=3))") + +sr_stoch_tdμ = solve_ECG(ops_tdμ, n_tdμ; scale = scale_tdμ, verbose = false) +ΔE_stoch_tdμ = sr_stoch_tdμ.ground_state - E_exact_tdμ +println(" Stochastic E₀ = $(round(sr_stoch_tdμ.ground_state, digits=4)) ΔE = $(round(ΔE_stoch_tdμ, sigdigits=3))") +println(" Exact E₀ = $E_exact_tdμ") + + +using Plots +n_s, E_s = convergence(sr_stoch) +plot(n_s, E_s; + label = "Stochastic", + xlabel = "Basis size", + ylabel = "Energy (Ha)", + title = "H⁻ convergence (n = $n)", +) +hline!([sr_var.ground_state]; label = "Variational (n=$n)", linestyle = :dash) +hline!([E_exact_Hm]; label = "Exact", linestyle = :dot, color = :black) \ No newline at end of file diff --git a/Project.toml b/Project.toml index a871bcc..92dc1fc 100644 --- a/Project.toml +++ b/Project.toml @@ -5,14 +5,18 @@ authors = ["Shuhei Ohno", "Martin Mikkelsen"] [deps] FewBodyHamiltonians = "3a126c26-e5d7-4a95-83c3-3b69f8a11ded" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +OptimKit = "77e91f04-9b3b-57a6-a776-40b61faaebe0" QuasiMonteCarlo = "8a4e6c94-4038-4cdc-81c3-7e6ffdb2a71b" SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" [compat] Aqua = "0.8.13" FewBodyHamiltonians = "0.0.2" +ForwardDiff = "1.3.2" LinearAlgebra = "1.7.3" +OptimKit = "0.4.2" QuasiMonteCarlo = "0.3.3" SpecialFunctions = "2.5.0" Test = "1.11.0" diff --git a/src/FewBodyECG.jl b/src/FewBodyECG.jl index 90e84e2..a965dfa 100644 --- a/src/FewBodyECG.jl +++ b/src/FewBodyECG.jl @@ -17,12 +17,15 @@ export build_hamiltonian_matrix, build_overlap_matrix, solve_generalized_eigenpr export ψ₀, SolverResults, convergence, correlation_function, ψ +export solve_ECG_variational + include("types.jl") include("coordinates.jl") include("matrix_elements.jl") include("hamiltonian.jl") include("sampling.jl") include("utils.jl") +include("variational.jl") end diff --git a/src/hamiltonian.jl b/src/hamiltonian.jl index 995aebe..9c1ba36 100644 --- a/src/hamiltonian.jl +++ b/src/hamiltonian.jl @@ -7,7 +7,8 @@ end function build_overlap_matrix(basis::BasisSet{<:GaussianBase}) n = length(basis.functions) - S = Matrix{Float64}(undef, n, n) + T = eltype(parent(first(basis.functions).A)) + S = Matrix{T}(undef, n, n) for i in 1:n, j in 1:i val = _compute_overlap_element(basis.functions[i], basis.functions[j]) S[i, j] = val @@ -18,7 +19,8 @@ end function _build_operator_matrix(basis::BasisSet{<:GaussianBase}, op::FewBodyHamiltonians.Operator) n = length(basis.functions) - H = Matrix{Float64}(undef, n, n) + T = eltype(parent(first(basis.functions).A)) + H = Matrix{T}(undef, n, n) for i in 1:n, j in 1:i val = _compute_matrix_element(basis.functions[i], basis.functions[j], op) H[i, j] = val @@ -29,7 +31,8 @@ end function build_hamiltonian_matrix(basis::BasisSet{<:GaussianBase}, operators::AbstractVector{<:FewBodyHamiltonians.Operator}) n = length(basis.functions) - H = zeros(Float64, n, n) + T = eltype(parent(first(basis.functions).A)) + H = zeros(T, n, n) for op in operators H .+= _build_operator_matrix(basis, op) end diff --git a/src/matrix_elements.jl b/src/matrix_elements.jl index 10c604d..497b9c8 100644 --- a/src/matrix_elements.jl +++ b/src/matrix_elements.jl @@ -7,7 +7,7 @@ Compute the matrix element ⟨bra|op|ket⟩ using analytic expressions. """ function _compute_matrix_element(bra::Rank0Gaussian, ket::Rank0Gaussian) - A, B = bra.A, ket.A + A, B = parent(bra.A), parent(ket.A) a, b = bra.s, ket.s S = A + B R = inv(S) @@ -74,7 +74,7 @@ function _compute_matrix_element(bra::Rank2Gaussian, ket::Rank2Gaussian) end function _compute_matrix_element(bra::Rank0Gaussian, ket::Rank0Gaussian, op::KineticOperator) - A, B = bra.A, ket.A + A, B = parent(bra.A), parent(ket.A) a, b = bra.s, ket.s K = op.K S = A + B @@ -178,7 +178,7 @@ function _compute_matrix_element(bra::Rank2Gaussian, ket::Rank2Gaussian, op::Kin end function _compute_matrix_element(bra::Rank0Gaussian, ket::Rank0Gaussian, op::CoulombOperator) - A, B = bra.A, ket.A + A, B = parent(bra.A), parent(ket.A) a, b = bra.s, ket.s w = op.w S = A + B diff --git a/src/variational.jl b/src/variational.jl new file mode 100644 index 0000000..38825ad --- /dev/null +++ b/src/variational.jl @@ -0,0 +1,308 @@ +using OptimKit +using LinearAlgebra +using QuasiMonteCarlo +using ForwardDiff + +function _chol_to_params(L::AbstractMatrix) + n = size(L, 1) + params = Float64[] + for j in 1:n + for i in j:n + push!(params, i == j ? log(L[i, j]) : L[i, j]) + end + end + return params +end + +function _params_to_matrix(θ::AbstractVector, n::Int) + T = eltype(θ) + L = zeros(T, n, n) + idx = 1 + for j in 1:n + for i in j:n + L[i, j] = (i == j) ? exp(θ[idx]) : θ[idx] + idx += 1 + end + end + return @assert(isposdef(Symmetric(L * L'))) +end + +function _encode_basis(basis::BasisSet{<:Rank0Gaussian}) + params = Float64[] + for g in basis.functions + C = cholesky(Symmetric(Matrix(g.A))) + append!(params, _chol_to_params(Matrix(C.L))) + end + return params +end + +# Decode a flat parameter vector back into a BasisSet{Rank0Gaussian} with +# zero shift vectors. +function _decode_basis(θ::AbstractVector, n_basis::Int, n_dim::Int) + T = eltype(θ) + n_chol = n_dim * (n_dim + 1) ÷ 2 + fns = Vector{Rank0Gaussian{T, Matrix{T}, Vector{T}}}(undef, n_basis) + for i in 1:n_basis + start = (i - 1) * n_chol + 1 + A = _params_to_matrix(θ[start:(start + n_chol - 1)], n_dim) + fns[i] = Rank0Gaussian(Matrix(A), zeros(T, n_dim)) + end + return BasisSet(fns) +end + +# --------------------------------------------------------------------------- +# Loss functions +# --------------------------------------------------------------------------- + +# Minimum generalised eigenvalue λ_min of (H, S). +# By the variational principle λ_min ≥ E₀ for any basis, so minimising +# over basis parameters converges to the exact ground-state energy. +function _energy_loss( + θ::AbstractVector, + n_basis::Int, + n_dim::Int, + operators::AbstractVector{<:FewBodyHamiltonians.Operator}; + regularization::Real = 1.0e-10 + ) + basis = _decode_basis(θ, n_basis, n_dim) + H = build_hamiltonian_matrix(basis, operators) + S = build_overlap_matrix(basis) + evals, _ = solve_generalized_eigenproblem(H, S; regularization = regularization) + return minimum(evals) +end + +# Tr(S⁻¹H) = sum of all generalised eigenvalues. +# For a basis that already approximates the ground state well this provides +# a smooth surrogate for the energy, but for a random or warm-started basis +# whose upper eigenvalues are large and positive the optimizer can reach a +# degenerate near-zero minimum by spreading the Gaussians out. Prefer +# loss_type = :energy unless you know the initial trace is already negative. +function _trace_loss( + θ::AbstractVector, + n_basis::Int, + n_dim::Int, + operators::AbstractVector{<:FewBodyHamiltonians.Operator}; + regularization::Real = 1.0e-10 + ) + basis = _decode_basis(θ, n_basis, n_dim) + H = build_hamiltonian_matrix(basis, operators) + S = build_overlap_matrix(basis) + S_reg = Symmetric(S + regularization * I) + return tr(S_reg \ H) +end + +# --------------------------------------------------------------------------- +# Main solver +# --------------------------------------------------------------------------- + +""" + solve_ECG_variational(operators, n; kwargs...) -> SolverResults + +Minimise a variational loss over the parameters of a `Rank0Gaussian` ECG +basis using any OptimKit.jl optimisation algorithm. + +Two loss types are available via the `loss_type` keyword: + +* **`:energy`** (default) — minimise the lowest generalised eigenvalue + `λ_min` of `Hc = λSc`. By the variational principle `λ_min ≥ E₀` + for any basis, so its global minimum is the exact ground-state energy. + This is the recommended choice and is equivalent to the standard + Rayleigh-Ritz variational principle applied to the Gaussian parameters. + Gradients are computed via ForwardDiff and the Hellmann-Feynman theorem + (differentiating through H and S only, not the eigensolver). + +* **`:trace`** — minimise `Tr(S⁻¹H)`, the sum of all generalised + eigenvalues. For a basis already close to the physical optimum this + can accelerate convergence, but for a randomly initialised basis + whose upper eigenvalues are large and positive the optimizer may find + a degenerate near-zero minimum (all Gaussians spread to infinity). + Prefer `:trace` only when warm-starting from a stochastic result. + +The `A` matrices of the basis functions are parameterised through their +Cholesky factors (log-diagonal) which keeps them positive-definite for +any parameter vector, giving a smooth unconstrained problem. Shift +vectors are fixed at zero, which is appropriate for ground states with +spherical symmetry. + +After optimisation the full generalised eigenvalue problem `Hc = λSc` +is solved once to produce the ground-state energy and eigenvectors, so +all `SolverResults`-based utilities (`ψ₀`, `correlation_function`, etc.) +work unchanged. + +# Arguments +- `operators` : `Vector{<:Operator}` — kinetic + Coulomb operators +- `n` : number of basis functions (default 50) + +# Keyword arguments +| keyword | default | description | +|:-----------------|:---------------|:------------| +| `loss_type` | `:energy` | `:energy` (λ_min) or `:trace` (Tr(S⁻¹H)) | +| `initial_basis` | `nothing` | warm-start `BasisSet{<:Rank0Gaussian}`; fresh QMC basis built when `nothing` | +| `scale` | `0.2` | length scale for QMC initialisation (ignored if `initial_basis` given) | +| `optimizer` | `nothing` | any OptimKit algorithm (e.g. `LBFGS()`, `ConjugateGradient()`); when `nothing` an L-BFGS is built from `max_iterations`, `gradient_tol`, and `verbose` | +| `max_iterations` | `500` | max iterations for the default optimizer (ignored if `optimizer` is given) | +| `gradient_tol` | `1e-6` | gradient-norm tolerance for the default optimizer (ignored if `optimizer` is given) | +| `regularization` | `1e-10` | Tikhonov shift added to S | +| `verbose` | `true` | print solver info messages; also sets verbosity of the default optimizer | + +# Example + +```julia +using FewBodyECG, OptimKit +masses = [1.0e15, 1.0, 1.0] +Λmat = Λ(masses) +J, U = _jacobi_transform(masses) +w_list = [[1,-1,0],[1,0,-1],[0,1,-1]] +w_raw = [U'*w for w in w_list] +ops = Operator[KineticOperator(Λmat); + [CoulombOperator(c,w) for (c,w) in zip([-1.,-1.,1.], w_raw)]...] + +# Default L-BFGS (convenience kwargs control it) +sr = solve_ECG_variational(ops, 30; scale=1.0, max_iterations=300, verbose=false) +println(sr.ground_state) + +# Pass any OptimKit algorithm directly +sr2 = solve_ECG_variational(ops, 30; scale=1.0, + optimizer=LBFGS(; maxiter=500, gradtol=1e-8, verbosity=2)) +println(sr2.ground_state) + +# Warm-start refinement with conjugate gradient +sr0 = solve_ECG(ops, 30; scale=1.0, verbose=false) +basis0 = BasisSet(Rank0Gaussian[sr0.basis_functions...]) +sr_cg = solve_ECG_variational(ops, 30; initial_basis=basis0, loss_type=:trace, + optimizer=ConjugateGradient(; maxiter=200, verbosity=0)) +println(sr_cg.ground_state) +``` +""" +function solve_ECG_variational( + operators::Vector{<:FewBodyHamiltonians.Operator}, + n::Int = 50; + loss_type::Symbol = :energy, + initial_basis::Union{BasisSet{<:Rank0Gaussian}, Nothing} = nothing, + scale::Real = 0.2, + optimizer = nothing, + max_iterations::Int = 500, + gradient_tol::Real = 1.0e-6, + regularization::Real = 1.0e-10, + verbose::Bool = true + ) + + loss_type in (:energy, :trace) || + throw(ArgumentError("loss_type must be :energy or :trace, got :$loss_type")) + + # Infer the Jacobi-coordinate dimension from the kinetic operator. + n_dim = size(first(op for op in operators if op isa KineticOperator).K, 1) + n_chol = n_dim * (n_dim + 1) ÷ 2 # free parameters per Gaussian + + # ---- build initial basis ------------------------------------------------ + if initial_basis !== nothing + length(initial_basis.functions) == n || throw(ArgumentError( + "initial_basis has $(length(initial_basis.functions)) functions, expected $n" + )) + basis_init = initial_basis + else + w_list = [op.w for op in operators if op isa CoulombOperator] + b1 = float(scale) + fns = Rank0Gaussian[] + for i in 1:n + bij = generate_bij(:quasirandom, i, length(w_list), b1) + A = _generate_A_matrix(bij, w_list) + push!(fns, Rank0Gaussian(A, zeros(n_dim))) + end + basis_init = BasisSet(fns) + end + + θ0 = _encode_basis(basis_init) + + # ---- combined value + ForwardDiff gradient (OptimKit interface) --------- + # The LBFGS line search probes regions that can produce degenerate A + # matrices; returning (Inf, zero-gradient) acts as an infinite-cost barrier. + function fg(θ::AbstractVector) + if loss_type === :energy + # Primal: solve Float64 eigenproblem for λ_min and eigenvector c. + local val::Float64, c::Vector{Float64} + try + basis_f64 = _decode_basis(θ, n, n_dim) + H = build_hamiltonian_matrix(basis_f64, operators) + S = build_overlap_matrix(basis_f64) + evals, evecs = solve_generalized_eigenproblem(H, S; regularization) + idx = argmin(evals) + val = evals[idx] + c = evecs[:, idx] + catch + return Inf, zeros(Float64, length(θ)) + end + isfinite(val) || return Inf, zeros(Float64, length(θ)) + # Gradient via Hellmann-Feynman: ∂λ/∂θ = cᵀ(∂H/∂θ − λ·∂S/∂θ)c + G = try + ForwardDiff.gradient(θ) do θ_ad + basis_ad = _decode_basis(θ_ad, n, n_dim) + H_ad = build_hamiltonian_matrix(basis_ad, operators) + S_ad = build_overlap_matrix(basis_ad) + dot(c, H_ad * c) - val * dot(c, S_ad * c) + end + catch + zeros(Float64, length(θ)) + end + return val, G + else # :trace + local val_t::Float64 + try + basis_f64 = _decode_basis(θ, n, n_dim) + H = build_hamiltonian_matrix(basis_f64, operators) + S = build_overlap_matrix(basis_f64) + v = tr((S + regularization * I) \ H) + val_t = isfinite(v) ? v : Inf + catch + return Inf, zeros(Float64, length(θ)) + end + isfinite(val_t) || return Inf, zeros(Float64, length(θ)) + G = try + ForwardDiff.gradient(θ) do θ_ad + basis_ad = _decode_basis(θ_ad, n, n_dim) + H_ad = build_hamiltonian_matrix(basis_ad, operators) + S_ad = build_overlap_matrix(basis_ad) + tr((S_ad + regularization * I) \ H_ad) + end + catch + zeros(Float64, length(θ)) + end + return val_t, G + end + end + + # ---- optimise ----------------------------------------------------------- + method = if optimizer !== nothing + optimizer + else + LBFGS(; maxiter = max_iterations, gradtol = float(gradient_tol), + verbosity = verbose ? 2 : 0) + end + + verbose && @info "Starting variational ECG optimisation" n_basis = n n_params = length(θ0) loss_type + + θ_opt, f_opt, _, _, _ = optimize(fg, θ0, method) + + verbose && @info "Optimisation done" loss = f_opt + + # ---- reconstruct basis and solve eigenproblem --------------------------- + basis_opt = _decode_basis(θ_opt, n, n_dim) + H_opt = build_hamiltonian_matrix(basis_opt, operators) + S_opt = build_overlap_matrix(basis_opt) + evals, evecs = solve_generalized_eigenproblem(H_opt, S_opt) + ground_state = minimum(evals) + + @info "Variational ECG complete" E₀ = ground_state n_basis = n + + return SolverResults( + Vector{GaussianBase}(basis_opt.functions), + n, + operators, + :variational, + HaltonSample(), # placeholder: no stochastic sampler is used + float(scale), + ground_state, + [ground_state], # single-point; no greedy build-up history + [evecs], + ) +end diff --git a/test/runtests.jl b/test/runtests.jl index c7c9f3b..77e882b 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -12,5 +12,6 @@ using FewBodyECG include("test_hydrogen.jl") include("test_utils.jl") include("test_types.jl") + include("test_variational.jl") end diff --git a/test/test_variational.jl b/test/test_variational.jl new file mode 100644 index 0000000..53d1de4 --- /dev/null +++ b/test/test_variational.jl @@ -0,0 +1,218 @@ +using Test +using LinearAlgebra +using FewBodyECG +import FewBodyECG: _jacobi_transform, _encode_basis, _decode_basis, _chol_to_params, _params_to_matrix + +# --------------------------------------------------------------------------- +# Shared 2-body (hydrogen) and 3-body (H⁻) operator fixtures +# --------------------------------------------------------------------------- + +function _hydrogen_ops() + masses = [1.0e15, 1.0] + _, U = _jacobi_transform(masses) + w = U' * [1.0, -1.0] + ops = Operator[KineticOperator(Λ(masses)); CoulombOperator(-1.0, w)] + return ops +end + +function _hminus_ops() + masses = [1.0e15, 1.0, 1.0] + _, U = _jacobi_transform(masses) + w_list = [[1, -1, 0], [1, 0, -1], [0, 1, -1]] + w_raw = [U' * Float64.(w) for w in w_list] + coeffs = [-1.0, -1.0, +1.0] + ops = Operator[KineticOperator(Λ(masses)); + [CoulombOperator(c, w) for (c, w) in zip(coeffs, w_raw)]...] + return ops +end + +# --------------------------------------------------------------------------- +# Cholesky parameterisation helpers +# --------------------------------------------------------------------------- + +@testset "_chol_to_params / _params_to_matrix round-trip" begin + for n in [1, 2, 3] + L = LowerTriangular(tril(rand(n, n)) + 2I) # positive diagonal + params = _chol_to_params(Matrix(L)) + A_reconstructed = _params_to_matrix(params, n) + A_original = Symmetric(L * L') + @test A_reconstructed ≈ A_original rtol = 1.0e-10 + end +end + +@testset "_encode_basis / _decode_basis round-trip" begin + # 2-body (n_dim=1): 1×1 A matrices + g1 = Rank0Gaussian([2.0;;], [0.0]) + g2 = Rank0Gaussian([5.0;;], [0.0]) + basis = BasisSet([g1, g2]) + + θ = _encode_basis(basis) + @test length(θ) == 2 # n_chol(1) = 1 per function, 2 functions + + basis2 = _decode_basis(θ, 2, 1) + @test length(basis2.functions) == 2 + for (orig, recon) in zip(basis.functions, basis2.functions) + @test Matrix(orig.A) ≈ Matrix(recon.A) rtol = 1.0e-10 + end +end + +@testset "_encode_basis / _decode_basis round-trip (2D)" begin + # 3-body (n_dim=2): 2×2 A matrices + A = [3.0 0.5; 0.5 2.0] + g = Rank0Gaussian(A, [0.0, 0.0]) + basis = BasisSet([g]) + + θ = _encode_basis(basis) + @test length(θ) == 3 # n_chol(2) = 3 + + basis2 = _decode_basis(θ, 1, 2) + @test Matrix(basis2.functions[1].A) ≈ A rtol = 1.0e-8 +end + +# --------------------------------------------------------------------------- +# solve_ECG_variational — argument validation +# --------------------------------------------------------------------------- + +@testset "solve_ECG_variational argument validation" begin + ops = _hydrogen_ops() + + @testset "Unknown loss_type throws" begin + @test_throws ArgumentError solve_ECG_variational( + ops, 3; loss_type = :bad, verbose = false + ) + end + + @testset "Mismatched initial_basis size throws" begin + sr = solve_ECG(ops, 5; verbose = false, scale = 1.0) + basis5 = BasisSet(Rank0Gaussian[sr.basis_functions...]) + @test_throws ArgumentError solve_ECG_variational( + ops, 3; initial_basis = basis5, verbose = false + ) + end +end + +# --------------------------------------------------------------------------- +# solve_ECG_variational — returned SolverResults structure +# --------------------------------------------------------------------------- + +@testset "solve_ECG_variational returns valid SolverResults" begin + ops = _hydrogen_ops() + sr = solve_ECG_variational(ops, 5; + scale = 1.0, max_iterations = 20, verbose = false + ) + + @test sr isa SolverResults + @test sr.n_basis == 5 + @test length(sr.basis_functions) == 5 + @test isfinite(sr.ground_state) + @test sr.ground_state < 0.0 # bound state + @test sr.method === :variational + @test length(sr.energies) == 1 + @test sr.energies[1] == sr.ground_state + @test length(sr.eigenvectors) == 1 + @test size(sr.eigenvectors[1]) == (5, 5) +end + +# --------------------------------------------------------------------------- +# solve_ECG_variational — both loss types run without error +# --------------------------------------------------------------------------- + +@testset "solve_ECG_variational loss_type = :energy" begin + ops = _hydrogen_ops() + sr = solve_ECG_variational(ops, 5; + loss_type = :energy, scale = 1.0, max_iterations = 20, verbose = false + ) + @test isfinite(sr.ground_state) + @test sr.ground_state < 0.0 +end + +@testset "solve_ECG_variational loss_type = :trace (warm start)" begin + ops = _hydrogen_ops() + # Warm-start from stochastic so the initial trace is already negative + sr_s = solve_ECG(ops, 5; scale = 1.0, verbose = false) + basis0 = BasisSet(Rank0Gaussian[sr_s.basis_functions...]) + sr = solve_ECG_variational(ops, 5; + loss_type = :trace, initial_basis = basis0, + max_iterations = 20, verbose = false + ) + @test isfinite(sr.ground_state) + @test sr.ground_state < 0.0 +end + +# --------------------------------------------------------------------------- +# solve_ECG_variational — variational principle +# --------------------------------------------------------------------------- + +@testset "solve_ECG_variational respects variational bound (hydrogen)" begin + ops = _hydrogen_ops() + E_exact = -0.5 # hydrogen 1s + + sr = solve_ECG_variational(ops, 10; + scale = 1.0, max_iterations = 100, verbose = false + ) + + # Variational principle: E₀ ≥ E_exact + @test sr.ground_state >= E_exact - 1.0e-6 + # With 10 functions, should get within 0.01 Ha of exact + @test sr.ground_state < E_exact + 0.01 +end + +@testset "solve_ECG_variational beats stochastic for hydrogen (same n)" begin + ops = _hydrogen_ops() + + sr_stoch = solve_ECG(ops, 8; scale = 1.0, verbose = false) + sr_var = solve_ECG_variational(ops, 8; + scale = 1.0, max_iterations = 150, verbose = false + ) + + # Optimised basis should be at least as good as the stochastic one + @test sr_var.ground_state <= sr_stoch.ground_state + 1.0e-6 +end + +# --------------------------------------------------------------------------- +# solve_ECG_variational — warm-start from stochastic result +# --------------------------------------------------------------------------- + +@testset "solve_ECG_variational warm-start improves stochastic result" begin + ops = _hminus_ops() + + sr_s = solve_ECG(ops, 8; scale = 1.0, verbose = false) + basis0 = BasisSet(Rank0Gaussian[sr_s.basis_functions...]) + + sr_v = solve_ECG_variational(ops, 8; + initial_basis = basis0, max_iterations = 100, verbose = false + ) + + # Variational principle holds + @test sr_v.ground_state >= -0.528 - 1.0e-4 + # Should not be worse than the starting point + @test sr_v.ground_state <= sr_s.ground_state + 1.0e-6 +end + +# --------------------------------------------------------------------------- +# Compatibility with downstream utilities +# --------------------------------------------------------------------------- + +@testset "ψ₀ works with variational SolverResults" begin + ops = _hydrogen_ops() + sr = solve_ECG_variational(ops, 5; + scale = 1.0, max_iterations = 20, verbose = false + ) + + r_vec = [0.5] # some point in Jacobi space + psi = ψ₀(r_vec, sr; state = 1) + @test isfinite(psi) +end + +@testset "correlation_function works with variational SolverResults" begin + ops = _hminus_ops() + sr = solve_ECG_variational(ops, 5; + scale = 1.0, max_iterations = 20, verbose = false + ) + + r_grid, rho = correlation_function(sr; npoints = 50) + @test length(r_grid) == 50 + @test length(rho) == 50 + @test all(isfinite, rho) + @test all(rho .>= 0.0) +end From 0175a8d12959f5978986823bcf3e0eb922cad8e5 Mon Sep 17 00:00:00 2001 From: MartinMikkelsen Date: Fri, 27 Feb 2026 18:07:12 +0100 Subject: [PATCH 2/8] added convergence loss --- Examples/Variational.jl | 27 ++++++++--- src/hamiltonian.jl | 97 ++++++++++++++++++++++++---------------- src/utils.jl | 13 ++++++ src/variational.jl | 38 ++++++++++++---- test/test_utils.jl | 9 ++-- test/test_variational.jl | 31 +++++++++++-- 6 files changed, 155 insertions(+), 60 deletions(-) diff --git a/Examples/Variational.jl b/Examples/Variational.jl index fd2c044..ece916c 100644 --- a/Examples/Variational.jl +++ b/Examples/Variational.jl @@ -117,12 +117,29 @@ println(" Exact E₀ = $E_exact_tdμ") using Plots +import FewBodyECG: convergence_history + +# --- Variational convergence (energy vs fg evaluations) --- +n_fg, E_fg = convergence_history(sr_var) +p1 = plot(n_fg, E_fg; + label = "Variational (cummin)", + xlabel = "fg evaluations", + ylabel = "Energy (Ha)", + title = "H⁻ variational convergence (n = $n)", + lw = 2, +) +hline!(p1, [E_exact_Hm]; label = "Exact", linestyle = :dot, color = :black, lw = 1) + +# --- Stochastic greedy convergence (energy vs basis size) --- n_s, E_s = convergence(sr_stoch) -plot(n_s, E_s; - label = "Stochastic", +p2 = plot(n_s, E_s; + label = "Stochastic greedy", xlabel = "Basis size", ylabel = "Energy (Ha)", - title = "H⁻ convergence (n = $n)", + title = "H⁻ stochastic convergence (n = $n)", + lw = 2, ) -hline!([sr_var.ground_state]; label = "Variational (n=$n)", linestyle = :dash) -hline!([E_exact_Hm]; label = "Exact", linestyle = :dot, color = :black) \ No newline at end of file +hline!(p2, [sr_var.ground_state]; label = "Variational (n=$n)", linestyle = :dash, lw = 1) +hline!(p2, [E_exact_Hm]; label = "Exact", linestyle = :dot, color = :black, lw = 1) + +plot(p1, p2; layout = (2, 1), size = (700, 600)) \ No newline at end of file diff --git a/src/hamiltonian.jl b/src/hamiltonian.jl index 9c1ba36..bf2c776 100644 --- a/src/hamiltonian.jl +++ b/src/hamiltonian.jl @@ -58,7 +58,6 @@ function solve_generalized_eigenproblem( cond_S = cond(S_sym) if cond_S > max_condition - @warn "Overlap matrix poorly conditioned (κ=$cond_S), adding regularization" if regularization == 0.0 regularization = maximum(abs.(diag(S_sym))) * 1.0e-10 end @@ -178,6 +177,11 @@ function solve_ECG( n_pairs = length(w_list) d = length(w_list[1]) + # Pre-allocate full matrices; fill one row/column per accepted function. + # S_full[j,j] doubles as a cache of self-overlaps for the independence check. + H_full = zeros(Float64, n, n) + S_full = zeros(Float64, n, n) + n_accepted = 0 n_rejected = 0 attempt = 0 @@ -190,61 +194,76 @@ function solve_ECG( s = generate_shift(method, attempt, d, scale; qmc_sampler = sampler) candidate = Rank0Gaussian(A, s) - # Check linear independence - if !isempty(basis_fns) - existing_basis = BasisSet{Rank0Gaussian}(basis_fns) - if !is_linearly_independent(candidate, existing_basis; threshold = threshold) + k = n_accepted # current accepted count + ki = k + 1 # index if this candidate is accepted + + # Compute new diagonal overlap (needed for independence check). + s_diag = _compute_matrix_element(candidate, candidate) + + # Compute new overlap column; check linear independence in the same pass. + s_col = Vector{Float64}(undef, k) + for j in 1:k + s_col[j] = _compute_matrix_element(candidate, basis_fns[j]) + end + if k > 0 + # S_full[j,j] holds the self-overlap of the j-th accepted function. + max_norm = maximum(j -> abs(s_col[j]) / sqrt(s_diag * S_full[j, j]), 1:k) + if max_norm > threshold n_rejected += 1 verbose && @warn "Rejected basis function $attempt (overlap > $threshold)" continue end end - push!(basis_fns, candidate) + # Compute new Hamiltonian column. + h_col = Vector{Float64}(undef, k) + for j in 1:k + h_col[j] = sum(_compute_matrix_element(candidate, basis_fns[j], op) for op in operators) + end + h_diag = sum(_compute_matrix_element(candidate, candidate, op) for op in operators) - local H, S, λs, Us - try - basis = BasisSet{Rank0Gaussian}(basis_fns) - H = build_hamiltonian_matrix(basis, operators) - S = build_overlap_matrix(basis) - - # Check for NaN/Inf BEFORE eigensolve - if any(!isfinite, H) - @warn "Hamiltonian contains NaN/Inf at step $(n_accepted + 1), rejecting basis function" - pop!(basis_fns) - n_rejected += 1 - continue - end + # Reject before touching the matrices if any element is non-finite. + if !isfinite(s_diag) || !isfinite(h_diag) || + (k > 0 && (!all(isfinite, s_col) || !all(isfinite, h_col))) + @warn "NaN/Inf in matrix elements at step $ki, rejecting basis function" + n_rejected += 1 + continue + end - if any(!isfinite, S) - @warn "Overlap contains NaN/Inf at step $(n_accepted + 1), rejecting basis function" - pop!(basis_fns) - n_rejected += 1 - continue - end + # Fill the new row/column into the pre-allocated matrices. + for j in 1:k + S_full[ki, j] = s_col[j] + S_full[j, ki] = s_col[j] + H_full[ki, j] = h_col[j] + H_full[j, ki] = h_col[j] + end + S_full[ki, ki] = s_diag + H_full[ki, ki] = h_diag - # Check condition number - cond_S = cond(Symmetric(S)) - if cond_S > max_condition - @warn "Overlap poorly conditioned (κ=$cond_S) at step $(n_accepted + 1), rejecting" - pop!(basis_fns) - n_rejected += 1 - continue - end + # Extract ki×ki submatrices (copy needed: eigensolver may alias internally). + H_k = H_full[1:ki, 1:ki] + S_k = S_full[1:ki, 1:ki] - # Try to solve - λs, Us = solve_generalized_eigenproblem(H, S; max_condition) + # Condition check on the overlap submatrix. + cond_S = cond(Symmetric(S_k)) + if cond_S > max_condition + @warn "Overlap poorly conditioned (κ=$cond_S) at step $ki, rejecting" + n_rejected += 1 + continue + end + local λs, Us + try + λs, Us = solve_generalized_eigenproblem(H_k, S_k; max_condition) catch e - @warn "Failed at step $(n_accepted + 1): $e" - pop!(basis_fns) # Remove problematic function + @warn "Failed at step $ki: $e" n_rejected += 1 continue end + push!(basis_fns, candidate) n_accepted += 1 E0 = minimum(λs) - push!(E_hist, E0) push!(vecs_list, Us) verbose && @info "Step $n_accepted" E₀ = E0 attempts = attempt rejected = n_rejected @@ -256,5 +275,5 @@ function solve_ECG( Emin = last(E_hist) @info "Optimization complete" E₀ = Emin n_basis = n_accepted - return SolverResults(basis_fns, n_accepted, operators, method, sampler, b₁, Emin, E_hist, vecs_list) + return SolverResults(basis_fns, n_accepted, operators, method, sampler, b₁, Emin, E_hist, vecs_list, E_hist) end diff --git a/src/utils.jl b/src/utils.jl index 3149100..32b6cc9 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -8,6 +8,11 @@ struct SolverResults ground_state::Float64 energies::Vector{Float64} eigenvectors::Vector{Matrix{Float64}} + # Per-call objective history (energy for :energy loss, trace for :trace loss). + # For solve_ECG this mirrors energies; for solve_ECG_variational it records + # the cumulative minimum at every primal fg evaluation so the curve is + # monotone and plottable via convergence_history(). + fg_history::Vector{Float64} end function ψ₀(r::AbstractVector, c::AbstractVector, basis_fns::Vector{<:GaussianBase}) @@ -26,6 +31,14 @@ function convergence(sr::SolverResults) return 1:sr.n_basis, sr.energies end +# Per-fg-evaluation objective history. For solve_ECG_variational with +# loss_type = :energy this is the energy at each primal solve, already +# reduced to a cumulative minimum so the curve is monotone. x-axis is +# the fg-call index, not the basis size. +function convergence_history(sr::SolverResults) + return 1:length(sr.fg_history), sr.fg_history +end + function correlation_function( sr::SolverResults; rmin::Real = 0.01, diff --git a/src/variational.jl b/src/variational.jl index 38825ad..2c32dfa 100644 --- a/src/variational.jl +++ b/src/variational.jl @@ -24,7 +24,7 @@ function _params_to_matrix(θ::AbstractVector, n::Int) idx += 1 end end - return @assert(isposdef(Symmetric(L * L'))) + return Symmetric(L * L') end function _encode_basis(basis::BasisSet{<:Rank0Gaussian}) @@ -32,20 +32,23 @@ function _encode_basis(basis::BasisSet{<:Rank0Gaussian}) for g in basis.functions C = cholesky(Symmetric(Matrix(g.A))) append!(params, _chol_to_params(Matrix(C.L))) + append!(params, Float64.(g.s)) # shift vector (unconstrained) end return params end -# Decode a flat parameter vector back into a BasisSet{Rank0Gaussian} with -# zero shift vectors. +# Decode a flat parameter vector back into a BasisSet{Rank0Gaussian}. +# Layout per Gaussian: [n_chol Cholesky params | n_dim shift params]. function _decode_basis(θ::AbstractVector, n_basis::Int, n_dim::Int) T = eltype(θ) n_chol = n_dim * (n_dim + 1) ÷ 2 + n_per = n_chol + n_dim fns = Vector{Rank0Gaussian{T, Matrix{T}, Vector{T}}}(undef, n_basis) for i in 1:n_basis - start = (i - 1) * n_chol + 1 + start = (i - 1) * n_per + 1 A = _params_to_matrix(θ[start:(start + n_chol - 1)], n_dim) - fns[i] = Rank0Gaussian(Matrix(A), zeros(T, n_dim)) + s = θ[(start + n_chol):(start + n_per - 1)] + fns[i] = Rank0Gaussian(Matrix(A), Vector(s)) end return BasisSet(fns) end @@ -192,7 +195,8 @@ function solve_ECG_variational( # Infer the Jacobi-coordinate dimension from the kinetic operator. n_dim = size(first(op for op in operators if op isa KineticOperator).K, 1) - n_chol = n_dim * (n_dim + 1) ÷ 2 # free parameters per Gaussian + n_chol = n_dim * (n_dim + 1) ÷ 2 # Cholesky params per Gaussian + n_per = n_chol + n_dim # total params per Gaussian (A + shift) # ---- build initial basis ------------------------------------------------ if initial_basis !== nothing @@ -214,6 +218,17 @@ function solve_ECG_variational( θ0 = _encode_basis(basis_init) + # Pre-allocate GradientConfig with a tuned chunk size. + # Grouping ~5 Gaussians per chunk gives ~10 passes for n=50 in 2D + # (vs the ForwardDiff default of 13), with larger gains for higher n_dim. + _chunk = min(n_per * 5, length(θ0)) + _grad_cfg = ForwardDiff.GradientConfig(nothing, θ0, ForwardDiff.Chunk(_chunk)) + + # Accumulates the objective value at every successful primal fg evaluation. + # Reduced to a cumulative minimum after optimisation so convergence_history() + # returns a monotone curve. + energy_log = Float64[] + # ---- combined value + ForwardDiff gradient (OptimKit interface) --------- # The LBFGS line search probes regions that can produce degenerate A # matrices; returning (Inf, zero-gradient) acts as an infinite-cost barrier. @@ -233,9 +248,10 @@ function solve_ECG_variational( return Inf, zeros(Float64, length(θ)) end isfinite(val) || return Inf, zeros(Float64, length(θ)) + push!(energy_log, val) # Gradient via Hellmann-Feynman: ∂λ/∂θ = cᵀ(∂H/∂θ − λ·∂S/∂θ)c G = try - ForwardDiff.gradient(θ) do θ_ad + ForwardDiff.gradient(θ, _grad_cfg, Val(false)) do θ_ad basis_ad = _decode_basis(θ_ad, n, n_dim) H_ad = build_hamiltonian_matrix(basis_ad, operators) S_ad = build_overlap_matrix(basis_ad) @@ -257,8 +273,9 @@ function solve_ECG_variational( return Inf, zeros(Float64, length(θ)) end isfinite(val_t) || return Inf, zeros(Float64, length(θ)) + push!(energy_log, val_t) G = try - ForwardDiff.gradient(θ) do θ_ad + ForwardDiff.gradient(θ, _grad_cfg, Val(false)) do θ_ad basis_ad = _decode_basis(θ_ad, n, n_dim) H_ad = build_hamiltonian_matrix(basis_ad, operators) S_ad = build_overlap_matrix(basis_ad) @@ -294,6 +311,10 @@ function solve_ECG_variational( @info "Variational ECG complete" E₀ = ground_state n_basis = n + # Reduce energy_log to a cumulative minimum so convergence_history() returns + # a monotone decreasing curve regardless of line-search noise. + fg_history = isempty(energy_log) ? Float64[] : accumulate(min, energy_log) + return SolverResults( Vector{GaussianBase}(basis_opt.functions), n, @@ -304,5 +325,6 @@ function solve_ECG_variational( ground_state, [ground_state], # single-point; no greedy build-up history [evecs], + fg_history, ) end diff --git a/test/test_utils.jl b/test/test_utils.jl index 3811122..bd5a396 100644 --- a/test/test_utils.jl +++ b/test/test_utils.jl @@ -45,7 +45,8 @@ function create_mock_solver_results(; scale, energies[end], energies, - eigenvectors + eigenvectors, + energies # fg_history mirrors energies for mock/stochastic results ) end @@ -78,8 +79,8 @@ end basis_fns = [Rank0Gaussian([1.0;;], [0.0])] ops = Operator[KineticOperator([0.5;;])] - sr_halton = SolverResults(basis_fns, 1, ops, :quasirandom, HaltonSample(), 1.0, -0.5, [-0.5], [ones(1, 1)]) - sr_sobol = SolverResults(basis_fns, 1, ops, :quasirandom, SobolSample(), 1.0, -0.5, [-0.5], [ones(1, 1)]) + sr_halton = SolverResults(basis_fns, 1, ops, :quasirandom, HaltonSample(), 1.0, -0.5, [-0.5], [ones(1, 1)], [-0.5]) + sr_sobol = SolverResults(basis_fns, 1, ops, :quasirandom, SobolSample(), 1.0, -0.5, [-0.5], [ones(1, 1)], [-0.5]) @test sr_halton.sampler isa HaltonSample @test sr_sobol.sampler isa SobolSample @@ -389,7 +390,7 @@ end sr = SolverResults( basis_fns, 1, ops, :quasirandom, HaltonSample(), - 1.0, -0.5, [-0.5], eigvecs + 1.0, -0.5, [-0.5], eigvecs, [-0.5] ) # All utilities should work diff --git a/test/test_variational.jl b/test/test_variational.jl index 53d1de4..4bb526d 100644 --- a/test/test_variational.jl +++ b/test/test_variational.jl @@ -41,32 +41,51 @@ end end @testset "_encode_basis / _decode_basis round-trip" begin - # 2-body (n_dim=1): 1×1 A matrices + # 2-body (n_dim=1): 1×1 A matrices, 1-D shift vectors + # n_per = n_chol(1) + n_dim(1) = 2 params per function g1 = Rank0Gaussian([2.0;;], [0.0]) g2 = Rank0Gaussian([5.0;;], [0.0]) basis = BasisSet([g1, g2]) θ = _encode_basis(basis) - @test length(θ) == 2 # n_chol(1) = 1 per function, 2 functions + @test length(θ) == 4 # 2 functions × (n_chol=1 + n_dim=1) basis2 = _decode_basis(θ, 2, 1) @test length(basis2.functions) == 2 for (orig, recon) in zip(basis.functions, basis2.functions) @test Matrix(orig.A) ≈ Matrix(recon.A) rtol = 1.0e-10 + @test recon.s ≈ orig.s rtol = 1.0e-10 end end @testset "_encode_basis / _decode_basis round-trip (2D)" begin - # 3-body (n_dim=2): 2×2 A matrices + # 3-body (n_dim=2): 2×2 A matrices, 2-D shift vectors + # n_per = n_chol(2) + n_dim(2) = 5 params per function A = [3.0 0.5; 0.5 2.0] g = Rank0Gaussian(A, [0.0, 0.0]) basis = BasisSet([g]) θ = _encode_basis(basis) - @test length(θ) == 3 # n_chol(2) = 3 + @test length(θ) == 5 # 1 function × (n_chol=3 + n_dim=2) basis2 = _decode_basis(θ, 1, 2) @test Matrix(basis2.functions[1].A) ≈ A rtol = 1.0e-8 + @test basis2.functions[1].s ≈ g.s rtol = 1.0e-10 +end + +@testset "_encode_basis / _decode_basis round-trip with non-zero shifts" begin + # Verify shift vectors are correctly preserved through the encode/decode cycle. + A = [3.0 0.5; 0.5 2.0] + s = [0.3, -0.1] + g = Rank0Gaussian(A, s) + basis = BasisSet([g]) + + θ = _encode_basis(basis) + @test length(θ) == 5 + + basis2 = _decode_basis(θ, 1, 2) + @test Matrix(basis2.functions[1].A) ≈ A rtol = 1.0e-8 + @test basis2.functions[1].s ≈ s rtol = 1.0e-10 end # --------------------------------------------------------------------------- @@ -111,6 +130,10 @@ end @test sr.energies[1] == sr.ground_state @test length(sr.eigenvectors) == 1 @test size(sr.eigenvectors[1]) == (5, 5) + # fg_history records cumulative-minimum energies from primal evaluations. + @test !isempty(sr.fg_history) + @test last(sr.fg_history) <= sr.ground_state + 1.0e-8 + @test issorted(sr.fg_history; rev = true) # monotone non-increasing end # --------------------------------------------------------------------------- From 73a7baef479bdfff1e8368d8e142b22586628236 Mon Sep 17 00:00:00 2001 From: MartinMikkelsen Date: Sat, 28 Feb 2026 09:19:28 +0100 Subject: [PATCH 3/8] updated matrix elements --- Examples/Variational.jl | 2 +- docs/Manifest.toml | 59 +++++++++++++++++------ docs/Project.toml | 1 + docs/make.jl | 1 + docs/src/API.md | 36 +++++++++++++- docs/src/examples.md | 72 ++++++++++++++++++++++++++++ docs/src/theory.md | 87 +++++++++++++++++++++++++++++++++ src/FewBodyECG.jl | 2 +- src/coordinates.jl | 23 +++++++++ src/hamiltonian.jl | 101 +++++++++++++++++++++++++-------------- src/matrix_elements.jl | 6 ++- src/types.jl | 51 ++++++++++++++++++++ src/utils.jl | 74 ++++++++++++++++++++++++++-- src/variational.jl | 14 ++++-- test/test_variational.jl | 39 ++++++++++++++- 15 files changed, 505 insertions(+), 63 deletions(-) diff --git a/Examples/Variational.jl b/Examples/Variational.jl index ece916c..fb8c7f7 100644 --- a/Examples/Variational.jl +++ b/Examples/Variational.jl @@ -122,7 +122,7 @@ import FewBodyECG: convergence_history # --- Variational convergence (energy vs fg evaluations) --- n_fg, E_fg = convergence_history(sr_var) p1 = plot(n_fg, E_fg; - label = "Variational (cummin)", + label = "Variational", xlabel = "fg evaluations", ylabel = "Energy (Ha)", title = "H⁻ variational convergence (n = $n)", diff --git a/docs/Manifest.toml b/docs/Manifest.toml index e183b2a..9bc62dd 100644 --- a/docs/Manifest.toml +++ b/docs/Manifest.toml @@ -2,7 +2,7 @@ julia_version = "1.12.5" manifest_format = "2.0" -project_hash = "c701bdbde2b173cbc9ce030ccf07a0026a315dfc" +project_hash = "3d39bcf94fc41feb6a82ef07b456752f08ca6ada" [[deps.ADTypes]] git-tree-sha1 = "f7304359109c768cf32dc5fa2d371565bb63b68a" @@ -211,9 +211,9 @@ version = "0.2.3" [[deps.ConcurrentUtilities]] deps = ["Serialization", "Sockets"] -git-tree-sha1 = "d9d26935a0bcffc87d2613ce14c527c99fc543fd" +git-tree-sha1 = "21d088c496ea22914fe80906eb5bce65755e5ec8" uuid = "f0e56b4a-5159-44fe-b623-3e5288b988bb" -version = "2.5.0" +version = "2.5.1" [[deps.ConstructionBase]] git-tree-sha1 = "b4b092499347b18a015186eae3042f72267106cb" @@ -337,9 +337,9 @@ version = "0.9.5" [[deps.Documenter]] deps = ["ANSIColoredPrinters", "AbstractTrees", "Base64", "CodecZlib", "Dates", "DocStringExtensions", "Downloads", "Git", "IOCapture", "InteractiveUtils", "JSON", "Logging", "Markdown", "MarkdownAST", "Pkg", "PrecompileTools", "REPL", "RegistryInstances", "SHA", "TOML", "Test", "Unicode"] -git-tree-sha1 = "b37458ae37d8bdb643d763451585cd8d0e5b4a9e" +git-tree-sha1 = "56e9c37b5e7c3b4f080ab1da18d72d5c290e184a" uuid = "e30172f5-a6a5-5a46-863b-614d45cd2de4" -version = "1.16.1" +version = "1.17.0" [[deps.Downloads]] deps = ["ArgTools", "FileWatching", "LibCURL", "NetworkOptions"] @@ -347,9 +347,9 @@ uuid = "f43a241f-c20a-4ad4-852c-f6b1247861c6" version = "1.7.0" [[deps.EnumX]] -git-tree-sha1 = "7bebc8aad6ee6217c78c5ddcf7ed289d65d0263e" +git-tree-sha1 = "c49898e8438c828577f04b92fc9368c388ac783c" uuid = "4e289a0a-7415-4d19-859d-a7e5c4648b56" -version = "1.0.6" +version = "1.0.7" [[deps.EpollShim_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl"] @@ -381,6 +381,12 @@ git-tree-sha1 = "01ba9d15e9eae375dc1eb9589df76b3572acd3f2" uuid = "b22a6f82-2f65-5046-a5b2-351ab43fb4e5" version = "8.0.1+0" +[[deps.FewBodyECG]] +deps = ["FewBodyHamiltonians", "ForwardDiff", "LinearAlgebra", "Logging", "OptimKit", "QuasiMonteCarlo", "SpecialFunctions"] +path = ".." +uuid = "083b1810-24a1-4a79-9a41-145bb2bb8ceb" +version = "1.0.5" + [[deps.FewBodyHamiltonians]] git-tree-sha1 = "3cf35661914b5eb9bce1f2a5ced704b4133ee1d4" uuid = "3a126c26-e5d7-4a95-83c3-3b69f8a11ded" @@ -555,6 +561,11 @@ git-tree-sha1 = "f923f9a774fcf3f5cb761bfa43aeadd689714813" uuid = "2e76f6c2-a576-52d4-95c1-20adfe4de566" version = "8.5.1+0" +[[deps.HashArrayMappedTries]] +git-tree-sha1 = "2eaa69a7cab70a52b9687c8bf950a5a93ec895ae" +uuid = "076d061b-32b6-4027-95e0-9a2c6f6d7e74" +version = "0.2.0" + [[deps.IOCapture]] deps = ["Logging", "Random"] git-tree-sha1 = "0ee181ec08df7d7c911901ea38baf16f755114dc" @@ -799,9 +810,9 @@ version = "0.1.3" [[deps.MbedTLS]] deps = ["Dates", "MbedTLS_jll", "MozillaCACerts_jll", "NetworkOptions", "Random", "Sockets"] -git-tree-sha1 = "c067a280ddc25f196b5e7df3877c6b226d390aaf" +git-tree-sha1 = "8785729fa736197687541f7053f6d8ab7fc44f92" uuid = "739be429-bea8-5141-9913-cc70e7f3736d" -version = "1.1.9" +version = "1.1.10" [[deps.MbedTLS_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl"] @@ -895,6 +906,12 @@ version = "1.13.3" [deps.Optim.weakdeps] MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee" +[[deps.OptimKit]] +deps = ["LinearAlgebra", "Printf", "ScopedValues", "VectorInterface"] +git-tree-sha1 = "5c92e3ab480969e80996587da5395f00224b9fdf" +uuid = "77e91f04-9b3b-57a6-a776-40b61faaebe0" +version = "0.4.2" + [[deps.Opus_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl"] git-tree-sha1 = "e2bb57a313a74b8104064b7efd01406c0a50d2ff" @@ -952,9 +969,9 @@ version = "1.4.4" [[deps.Plots]] deps = ["Base64", "Contour", "Dates", "Downloads", "FFMPEG", "FixedPointNumbers", "GR", "JLFzf", "JSON", "LaTeXStrings", "Latexify", "LinearAlgebra", "Measures", "NaNMath", "Pkg", "PlotThemes", "PlotUtils", "PrecompileTools", "Printf", "REPL", "Random", "RecipesBase", "RecipesPipeline", "Reexport", "RelocatableFolders", "Requires", "Scratch", "Showoff", "SparseArrays", "Statistics", "StatsBase", "TOML", "UUIDs", "UnicodeFun", "Unzip"] -git-tree-sha1 = "1cc8ad0762e59e713ee3ef28f9b78b2c9f4ca078" +git-tree-sha1 = "cb20a4eacda080e517e4deb9cfb6c7c518131265" uuid = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" -version = "1.41.5" +version = "1.41.6" [deps.Plots.extensions] FileIOExt = "FileIO" @@ -984,9 +1001,9 @@ version = "1.3.3" [[deps.Preferences]] deps = ["TOML"] -git-tree-sha1 = "522f093a29b31a93e34eaea17ba055d850edea28" +git-tree-sha1 = "8b770b60760d4451834fe79dd483e318eee709c4" uuid = "21216c6a-2e73-6563-6e65-726566657250" -version = "1.5.1" +version = "1.5.2" [[deps.Primes]] deps = ["IntegerMathUtils"] @@ -1000,9 +1017,9 @@ uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7" version = "1.11.0" [[deps.PtrArrays]] -git-tree-sha1 = "1d36ef11a9aaf1e8b74dacc6a731dd1de8fd493d" +git-tree-sha1 = "4fbbafbc6251b883f4d2705356f3641f3652a7fe" uuid = "43287f4e-b6f4-7ad1-bb20-aadabca52c3d" -version = "1.3.0" +version = "1.4.0" [[deps.Qt6Base_jll]] deps = ["Artifacts", "CompilerSupportLibraries_jll", "Fontconfig_jll", "Glib_jll", "JLLWrappers", "Libdl", "Libglvnd_jll", "OpenSSL_jll", "Vulkan_Loader_jll", "Xorg_libSM_jll", "Xorg_libXext_jll", "Xorg_libXrender_jll", "Xorg_libxcb_jll", "Xorg_xcb_util_cursor_jll", "Xorg_xcb_util_image_jll", "Xorg_xcb_util_keysyms_jll", "Xorg_xcb_util_renderutil_jll", "Xorg_xcb_util_wm_jll", "Zlib_jll", "libinput_jll", "xkbcommon_jll"] @@ -1089,6 +1106,12 @@ version = "1.3.1" uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce" version = "0.7.0" +[[deps.ScopedValues]] +deps = ["HashArrayMappedTries", "Logging"] +git-tree-sha1 = "c3b2323466378a2ba15bea4b2f73b081e022f473" +uuid = "7e506255-f358-4e82-b7e4-beb19740aa63" +version = "1.5.0" + [[deps.Scratch]] deps = ["Dates"] git-tree-sha1 = "9b81b8393e50b7d4e6d0a9f14e192294d3b7c109" @@ -1256,6 +1279,12 @@ git-tree-sha1 = "ca0969166a028236229f63514992fc073799bb78" uuid = "41fe7b60-77ed-43a1-b4f0-825fd5a5650d" version = "0.2.0" +[[deps.VectorInterface]] +deps = ["LinearAlgebra"] +git-tree-sha1 = "9166406dedd38c111a6574e9814be83d267f8aec" +uuid = "409d34a3-91d5-4945-b6ec-7529ddf182d8" +version = "0.5.0" + [[deps.Vulkan_Loader_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl", "Wayland_jll", "Xorg_libX11_jll", "Xorg_libXrandr_jll", "xkbcommon_jll"] git-tree-sha1 = "2f0486047a07670caad3a81a075d2e518acc5c59" diff --git a/docs/Project.toml b/docs/Project.toml index 17b7dec..2e161fa 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -1,5 +1,6 @@ [deps] Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" +FewBodyECG = "083b1810-24a1-4a79-9a41-145bb2bb8ceb" FewBodyHamiltonians = "3a126c26-e5d7-4a95-83c3-3b69f8a11ded" Optim = "429524aa-4258-5aef-a3af-852621145aeb" Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" diff --git a/docs/make.jl b/docs/make.jl index d54ae9d..d7a1be9 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -2,6 +2,7 @@ using Documenter, FewBodyECG makedocs( build = "build", + modules = [FewBodyECG], sitename = "FewBodyECG.jl", pages = [ "Home" => "index.md", diff --git a/docs/src/API.md b/docs/src/API.md index d6a7a7c..2e9efab 100644 --- a/docs/src/API.md +++ b/docs/src/API.md @@ -1 +1,35 @@ -# API +# API Reference + +## Solvers + +```@docs +solve_ECG +solve_ECG_variational +``` + +## Operators + +```@docs +KineticOperator +CoulombOperator +``` + +## Results and utilities + +```@docs +SolverResults +convergence +convergence_history +correlation_function +ψ₀ +``` + +## Coordinates and basis + +```@docs +Λ +_jacobi_transform +GaussianBase +Rank0Gaussian +BasisSet +``` diff --git a/docs/src/examples.md b/docs/src/examples.md index 226d16b..dc17b85 100644 --- a/docs/src/examples.md +++ b/docs/src/examples.md @@ -32,4 +32,76 @@ E = -0.527751016523 n, E = convergence(result) plot(n, E) +``` + +## Variational optimisation with `solve_ECG_variational` + +The stochastic solver above builds the basis greedily from random samples. +`solve_ECG_variational` instead treats all Gaussian parameters (the `A` +matrices encoded through their Cholesky factors, and the shift vectors `s`) +as continuous variables and minimises the ground-state energy directly with +L-BFGS via [OptimKit.jl](https://github.com/Jutho/OptimKit.jl). + +### Fresh start + +```@example example_var_fresh +using FewBodyECG +using LinearAlgebra +using Plots + +masses = [1.0e15, 1.0, 1.0] # H⁻: fixed nucleus + 2 electrons + +Λmat = Λ(masses) +kin = KineticOperator(Λmat) +J, U = _jacobi_transform(masses) + +w_list = [[1, -1, 0], [1, 0, -1], [0, 1, -1]] +w_raw = [U' * w for w in w_list] +coeffs = [-1.0, -1.0, +1.0] + +ops = Operator[kin; [CoulombOperator(c, w) for (c, w) in zip(coeffs, w_raw)]...] + +sr = solve_ECG_variational(ops, 20; scale = 1.0, max_iterations = 200, verbose = false) + +E_exact = -0.527751016523 +println("Variational E₀ = ", round(sr.ground_state, digits=8), + " ΔE = ", round(sr.ground_state - E_exact, sigdigits=3)) + +xs, ys = convergence_history(sr) +plot(xs, ys; xlabel="fg evaluations", ylabel="Energy (Ha)", + label="Variational", lw=2) +hline!([E_exact]; label="Exact", linestyle=:dot, color=:black) +``` + +### Warm start from stochastic result + +When the stochastic solver has already found a reasonable basis, refining it +with `loss_type = :trace` (minimise `Tr(S⁻¹H)`) often converges faster than a +fresh L-BFGS run. + +```@example example_var_warm +using FewBodyECG +using LinearAlgebra + +masses = [1.0e15, 1.0, 1.0] +Λmat = Λ(masses); kin = KineticOperator(Λmat) +J, U = _jacobi_transform(masses) +w_raw = [U' * w for w in [[1,-1,0],[1,0,-1],[0,1,-1]]] +ops = Operator[kin; [CoulombOperator(c,w) for (c,w) in zip([-1.,-1.,1.], w_raw)]...] + +sr_stoch = solve_ECG(ops, 20; scale = 1.0, verbose = false) +basis0 = BasisSet(Rank0Gaussian[sr_stoch.basis_functions...]) + +sr_warm = solve_ECG_variational(ops, 20; + initial_basis = basis0, + loss_type = :trace, + max_iterations = 200, + verbose = false, +) + +E_exact = -0.527751016523 +println("Stochastic E₀ = ", round(sr_stoch.ground_state, digits=8), + " ΔE = ", round(sr_stoch.ground_state - E_exact, sigdigits=3)) +println("Warm-start E₀ = ", round(sr_warm.ground_state, digits=8), + " ΔE = ", round(sr_warm.ground_state - E_exact, sigdigits=3)) ``` \ No newline at end of file diff --git a/docs/src/theory.md b/docs/src/theory.md index 24aecf1..bc5a4d9 100644 --- a/docs/src/theory.md +++ b/docs/src/theory.md @@ -72,3 +72,90 @@ and ```math \vec{r}^T A \vec{r} + \vec{s}^T \vec{r} = \sum_{i,j} \vec{r}_i \cdot A_{ij}\vec{r}_j + \sum_i \vec{s}_i \cdot \vec{r}_i. ``` + +## Stochastic basis construction + +The simplest strategy for choosing the Gaussian parameters is **stochastic +greedy search** (`solve_ECG`): candidate basis functions are generated from +quasi-random sequences (Halton, Sobol) and accepted one at a time if they +reduce the lowest eigenvalue. This is fast and robust, but the quality of the +final basis depends on the sampling distribution and the number of accepted +functions. + +## Variational parameter optimisation + +A more systematic approach is to treat all parameters of all basis functions +simultaneously as a continuous optimisation problem (`solve_ECG_variational`). + +### The variational principle + +By the Rayleigh-Ritz variational principle, the lowest generalised eigenvalue +$\lambda_{\min}$ of $Hc = \lambda S c$ satisfies + +```math +\lambda_{\min} \;\geq\; E_0, +``` + +where $E_0$ is the exact ground-state energy. Equality holds when the +parameter space is large enough to contain the true ground state. Therefore +minimising $\lambda_{\min}$ over the Gaussian parameters is a rigorous +upper-bound approach: the optimum is approached monotonically from above. + +### Parameterisation + +Every $n_{\text{dim}} \times n_{\text{dim}}$ positive-definite matrix $A$ is +written as $A = L L^T$ where $L$ is lower-triangular with positive diagonal. +Rather than optimising $L$ directly (which requires inequality constraints), +the diagonal entries are reparameterised as $L_{ii} = \exp(\theta_{ii})$, so +that the full parameter vector $\theta$ is unconstrained: + +```math +L_{ij}(\theta) = +\begin{cases} +e^{\theta_{ij}} & i = j, \\ +\theta_{ij} & i > j. +\end{cases} +``` + +The shift vector $s \in \mathbb{R}^{n_{\text{dim}}}$ is also included in +$\theta$ without any transformation. The complete parameter vector for a basis +of $n$ Gaussians therefore has dimension + +```math +|\theta| = n \times \left(\frac{n_{\text{dim}}(n_{\text{dim}}+1)}{2} + n_{\text{dim}}\right). +``` + +### Gradient computation (Hellmann-Feynman) + +Computing the gradient $\nabla_\theta \lambda_{\min}$ by automatic +differentiation through the eigensolver is expensive and numerically fragile. +Instead, the **Hellmann-Feynman theorem** is used: if $c$ is the eigenvector +corresponding to $\lambda_{\min}$, then + +```math +\frac{\partial \lambda_{\min}}{\partial \theta_k} += c^T \!\left(\frac{\partial H}{\partial \theta_k} - \lambda_{\min}\,\frac{\partial S}{\partial \theta_k}\right) c. +``` + +The eigenproblem is solved in `Float64` to obtain $\lambda_{\min}$ and $c$. +ForwardDiff then differentiates only through the matrix-build steps +($H(\theta)$ and $S(\theta)$), which avoids differentiating through any +eigenvalue decomposition. The chunk size of ForwardDiff is tuned so that each +Gaussian contributes approximately five dual-number columns per pass. + +### Optimisation + +The combined value-and-gradient function is passed to the L-BFGS +implementation provided by [OptimKit.jl](https://github.com/Jutho/OptimKit.jl). +Regions where the overlap matrix $S$ is near-singular (degenerate Gaussians) +are handled by returning $(+\infty, \mathbf{0})$ from the objective, which +acts as an infinite-cost barrier that the line search naturally avoids. + +### Trace loss (warm-start refinement) + +An alternative surrogate loss is $\operatorname{Tr}(S^{-1}H)$, the sum of +**all** generalised eigenvalues. This is differentiable everywhere without +requiring an eigensolver, but is only useful when the initial basis is already +close to the physical ground state (e.g. after a stochastic run), because for +random initialisations the upper eigenvalues are large and the optimizer may +find a trivial degenerate minimum. diff --git a/src/FewBodyECG.jl b/src/FewBodyECG.jl index a965dfa..9a48840 100644 --- a/src/FewBodyECG.jl +++ b/src/FewBodyECG.jl @@ -15,7 +15,7 @@ export Operator export build_hamiltonian_matrix, build_overlap_matrix, solve_generalized_eigenproblem, solve_ECG, convergence -export ψ₀, SolverResults, convergence, correlation_function, ψ +export ψ₀, SolverResults, convergence, convergence_history, correlation_function, ψ export solve_ECG_variational diff --git a/src/coordinates.jl b/src/coordinates.jl index e4948e0..5835347 100644 --- a/src/coordinates.jl +++ b/src/coordinates.jl @@ -1,3 +1,16 @@ +""" + _jacobi_transform(masses) -> (J, U) + +Compute the Jacobi coordinate transformation matrix `J` and its pseudo-inverse `U` +for a system with the given particle `masses`. + +Returns `(J, U)` where: +- `J` is the ``(N-1) \\times N`` matrix mapping particle coordinates to Jacobi + relative coordinates (centre-of-mass motion is factored out). +- `U = \\operatorname{pinv}(J)` is the ``N \\times (N-1)`` back-transformation. + +The weight vectors for `CoulombOperator` are constructed as `U' * charge_vector`. +""" function _jacobi_transform(masses::Vector{Float64})::Tuple{Matrix{Float64}, Matrix{Float64}} N = length(masses) @assert N ≥ 2 "At least two masses are required for Jacobi transformation." @@ -22,6 +35,16 @@ function _jacobi_transform(masses::Vector{Float64})::Tuple{Matrix{Float64}, Matr return J, U end +""" + Λ(masses) -> Symmetric matrix + +Compute the kinetic-energy matrix in Jacobi coordinates for a system with the +given particle `masses` (in atomic units). + +Returns the symmetric matrix ``\\Lambda = J M^{-1} J^T / 2``, where ``J`` is +the Jacobi transformation matrix and ``M = \\operatorname{diag}(m_i)``. +Pass the result directly to `KineticOperator`. +""" function Λ(masses::Vector{<:Real}) J, _ = _jacobi_transform(masses) Minv = Diagonal(0.5 ./ masses) diff --git a/src/hamiltonian.jl b/src/hamiltonian.jl index bf2c776..6997d05 100644 --- a/src/hamiltonian.jl +++ b/src/hamiltonian.jl @@ -64,58 +64,38 @@ function solve_generalized_eigenproblem( end if regularization > 0 - S_sym = S_sym + regularization * I + S_sym = Symmetric(Matrix(S_sym) + regularization * I) end if !isposdef(S_sym) @warn "Overlap matrix not positive definite, adding regularization" ε = maximum(abs.(diag(S_sym))) * 1.0e-8 - S_sym = S_sym + ε * I + S_sym = Symmetric(Matrix(S_sym) + ε * I) if !isposdef(S_sym) error("Overlap matrix not positive definite even after regularization") end end - # Cholesky decomposition with error handling - local F + # Solve the generalised symmetric eigenvalue problem H c = λ S c via + # LAPACK's divide-and-conquer driver (dsygvd). This is more reliable than + # manually factorising S and back-transforming, and returns eigenvectors + # normalised so that vᵀ S v = I. + local evals, vecs try - F = cholesky(S_sym) + F = eigen(H_sym, S_sym) + evals = real.(F.values) + vecs = real.(F.vectors) catch e - @error "Cholesky decomposition failed" exception = e - @error "Overlap matrix info" condition = cond(S_sym) min_eigval = minimum(eigvals(S_sym)) + @error "Generalised eigenvalue decomposition failed" exception = e rethrow(e) end - L = F.L - - # Transform to standard eigenvalue problem - A = (L \ Matrix(H_sym)) / L' - A_sym = Symmetric((A + A') / 2) - - # Check for NaN/Inf after transformation - if any(!isfinite, A_sym) - error("Transformed matrix contains NaN or Inf after Cholesky transformation") - end - - # Solve standard eigenvalue problem - local evals, evecs - try - evals, evecs = eigen(A_sym) - catch e - @error "Eigenvalue decomposition failed" exception = e - @error "Transformed matrix info" condition = cond(A_sym) - rethrow(e) + if any(!isfinite, evals) || any(!isfinite, vecs) + error("Eigenvalues or eigenvectors contain NaN or Inf") end - # Transform eigenvectors back - vecs = L' \ evecs - - # Ensure real - evals_real = real.(evals) - vecs_real = real.(vecs) - - return evals_real, vecs_real + return evals, vecs end function normalized_overlap(A::GaussianBase, B::GaussianBase) @@ -156,6 +136,46 @@ function default_scale(masses::Vector{<:Real}) return 1 / sqrt(μ) end +""" + solve_ECG(operators, n=50; kwargs...) -> SolverResults + +Build an ECG basis of `n` `Rank0Gaussian` functions using **stochastic greedy +search** and return the ground-state energy. + +Candidate Gaussians are generated from a quasi-random sequence (Halton by +default). Each candidate is accepted if it is linearly independent from the +existing basis (normalised overlap < `threshold`) and does not make the overlap +matrix ill-conditioned. The ground-state energy after each accepted function +is stored in `SolverResults.energies`. + +# Arguments +- `operators` : `Vector{<:Operator}` — kinetic + Coulomb operators (see [`KineticOperator`](@ref), [`CoulombOperator`](@ref)). +- `n` : target number of basis functions (default 50). + +# Keyword arguments +| keyword | default | description | +|:----------------|:---------------|:------------| +| `sampler` | `HaltonSample()` | QuasiMonteCarlo sampler for generating candidates | +| `method` | `:quasirandom` | `:quasirandom` or `:random` | +| `scale` | `0.2` | characteristic Gaussian width (a.u.) | +| `threshold` | `0.95` | normalised overlap above which a candidate is rejected | +| `max_attempts` | `10n` | maximum number of candidate draws | +| `max_condition` | `1e12` | maximum condition number of the overlap matrix | +| `verbose` | `true` | print per-step info messages | + +# Example + +```julia +using FewBodyECG +masses = [1.0e15, 1.0] # hydrogen atom (fixed nucleus) +Λmat = Λ(masses) +_, U = _jacobi_transform(masses) +w = U' * [1.0, -1.0] +ops = Operator[KineticOperator(Λmat); CoulombOperator(-1.0, w)] +sr = solve_ECG(ops, 30; scale=1.0, verbose=false) +println(sr.ground_state) # ≈ -0.5 Ha +``` +""" function solve_ECG( operators::Vector{<:FewBodyHamiltonians.Operator}, n::Int = 50; @@ -261,9 +281,20 @@ function solve_ECG( continue end + E0 = minimum(λs) + + # Variational principle: adding any linearly independent function to the + # basis cannot raise the ground-state energy. If it does, the candidate + # is numerically near-degenerate with the existing basis (not caught by + # the overlap / condition-number checks above), so reject it. + if n_accepted > 0 && E0 > E_hist[end] + 1.0e-10 + @warn "Candidate raises energy at step $ki, rejecting" ΔE = E0 - E_hist[end] + n_rejected += 1 + continue + end + push!(basis_fns, candidate) n_accepted += 1 - E0 = minimum(λs) push!(E_hist, E0) push!(vecs_list, Us) verbose && @info "Step $n_accepted" E₀ = E0 attempts = attempt rejected = n_rejected diff --git a/src/matrix_elements.jl b/src/matrix_elements.jl index 497b9c8..1a2e454 100644 --- a/src/matrix_elements.jl +++ b/src/matrix_elements.jl @@ -186,7 +186,11 @@ function _compute_matrix_element(bra::Rank0Gaussian, ket::Rank0Gaussian, op::Cou M = _compute_matrix_element(bra, ket) β = 1 / (w' * R * w) q = 0.5 * (w' * R * (a + b)) - f = abs(q) < 1.0e-12 ? (2 * sqrt(β / π)) : (erf(sqrt(β) * q) / q) + # Use the limiting form 2√(β/π) when the scaled argument x = √β·q is + # small so that erf(x)/q = √β·erf(x)/x → 2√(β/π) accurately. + # Thresholding on x (not q alone) correctly handles all β values. + x = sqrt(β) * q + f = abs(x) < 1.0e-7 ? (2 * sqrt(β / π)) : (erf(x) / q) return op.coefficient * f * M end diff --git a/src/types.jl b/src/types.jl index 8375955..88004be 100644 --- a/src/types.jl +++ b/src/types.jl @@ -1,7 +1,24 @@ using FewBodyHamiltonians +""" + GaussianBase + +Abstract supertype for all explicitly correlated Gaussian basis functions. +Concrete subtypes differ by the rank of the polynomial prefactor: +`Rank0Gaussian` (plain Gaussian), `Rank1Gaussian` (linear prefactor), +`Rank2Gaussian` (quadratic prefactor). +""" abstract type GaussianBase end +""" + Rank0Gaussian(A, s) + +Basis function ``g(\\mathbf{r}) = \\exp(-\\mathbf{r}^T A\\,\\mathbf{r} + \\mathbf{s}^T\\mathbf{r})``. + +# Fields +- `A` : symmetric positive-definite ``n_{\\text{dim}} \\times n_{\\text{dim}}`` matrix controlling the Gaussian width and inter-particle correlations. +- `s` : shift vector ``\\mathbf{s} \\in \\mathbb{R}^{n_{\\text{dim}}}``; controls the location of the Gaussian maximum. +""" struct Rank0Gaussian{T <: Real, M <: AbstractMatrix{T}, V <: AbstractVector{T}} <: GaussianBase A::Symmetric{T, M} s::V @@ -37,14 +54,48 @@ struct Rank2Gaussian{T <: Real, M <: AbstractMatrix{T}, V <: AbstractVector{T}} end end +""" + BasisSet(functions) + +A collection of `GaussianBase` functions that form the variational basis. + +# Fields +- `functions` : `Vector{G}` of basis functions, all of the same concrete `GaussianBase` subtype `G`. +""" struct BasisSet{G <: GaussianBase} functions::Vector{G} end +""" + KineticOperator(K) + KineticOperator(masses) + +Kinetic-energy operator in Jacobi coordinates. + +When constructed from a mass vector the Jacobi-transformed kinetic-energy +matrix ``\\Lambda = J M^{-1} J^T / 2`` is computed automatically via [`Λ`](@ref). + +# Fields +- `K` : symmetric ``n_{\\text{dim}} \\times n_{\\text{dim}}`` kinetic-energy matrix (``\\Lambda``). +""" struct KineticOperator{T <: Real} <: FewBodyHamiltonians.KineticTerm K::AbstractMatrix{T} end +""" + CoulombOperator(coefficient, w) + +Two-body Coulomb (``1/r_{ij}``) interaction operator. + +The inter-particle distance is ``|w^T \\mathbf{r}|`` where `w` is a weight +vector in Jacobi coordinates selecting the pair ``(i,j)``. Construct `w` +by transforming the charge-difference vector with the inverse Jacobi matrix: +`w = U' * charge_vector`. + +# Fields +- `coefficient` : coupling constant (e.g. ``q_i q_j``; negative for attraction). +- `w` : weight vector in Jacobi coordinates. +""" struct CoulombOperator{T <: Real} <: FewBodyHamiltonians.PotentialTerm coefficient::T w::AbstractVector{T} diff --git a/src/utils.jl b/src/utils.jl index 32b6cc9..47267a8 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -1,3 +1,25 @@ +""" + SolverResults + +Output of [`solve_ECG`](@ref) and [`solve_ECG_variational`](@ref). + +# Fields +| field | type | description | +|:------------------|:--------------------------------------|:------------| +| `basis_functions` | `Vector{GaussianBase}` | optimised basis | +| `n_basis` | `Int` | number of accepted/optimised functions | +| `operators` | `Vector{Operator}` | kinetic + Coulomb operators passed to the solver | +| `method` | `Symbol` | `:quasirandom`, `:random`, or `:variational` | +| `sampler` | `DeterministicSamplingAlgorithm` | QMC sampler used (placeholder for variational results) | +| `length_scale` | `Float64` | Gaussian width scale | +| `ground_state` | `Float64` | lowest eigenvalue (ground-state energy in Hartree) | +| `energies` | `Vector{Float64}` | energy at each greedy step (stochastic) or `[ground_state]` (variational) | +| `eigenvectors` | `Vector{Matrix{Float64}}` | eigenvector matrices at each step | +| `fg_history` | `Vector{Float64}` | cumulative-minimum energy per objective call (variational) or mirrors `energies` (stochastic) | + +Use [`convergence`](@ref), [`convergence_history`](@ref), [`correlation_function`](@ref), +and [`ψ₀`](@ref) to analyse the result. +""" struct SolverResults basis_functions::Vector{GaussianBase} n_basis::Int @@ -15,6 +37,18 @@ struct SolverResults fg_history::Vector{Float64} end +""" + ψ₀(r, c, basis_fns) + ψ₀(r, sr; state=1) + +Evaluate the ground-state wavefunction at Jacobi-coordinate point `r`. + +The wavefunction is the linear combination +``\\psi_0(\\mathbf{r}) = \\sum_i c_i \\exp(-\\mathbf{r}^T A_i \\mathbf{r} + \\mathbf{s}_i^T \\mathbf{r})``. + +When called with a [`SolverResults`](@ref), the eigenvector for the requested +`state` (default 1, i.e. the ground state) is used automatically. +""" function ψ₀(r::AbstractVector, c::AbstractVector, basis_fns::Vector{<:GaussianBase}) return sum( c[i] * exp(-r' * basis_fns[i].A * r + basis_fns[i].s' * r) @@ -27,18 +61,50 @@ function ψ₀(r::AbstractVector, sr::SolverResults; state::Int = 1) return ψ₀(r, c, sr.basis_functions) end +""" + convergence(sr::SolverResults) -> (indices, energies) + +Return the greedy build-up convergence curve from a stochastic [`solve_ECG`](@ref) run. + +Returns `(1:n_basis, sr.energies)`: the energy after each basis function was +added. For variational results `energies` contains only one entry +`[ground_state]`; use [`convergence_history`](@ref) instead. +""" function convergence(sr::SolverResults) return 1:sr.n_basis, sr.energies end -# Per-fg-evaluation objective history. For solve_ECG_variational with -# loss_type = :energy this is the energy at each primal solve, already -# reduced to a cumulative minimum so the curve is monotone. x-axis is -# the fg-call index, not the basis size. +""" + convergence_history(sr::SolverResults) -> (indices, energies) + +Return the per-objective-call convergence history. + +For [`solve_ECG_variational`](@ref) results this is the cumulative-minimum +energy at every primal `fg` evaluation, giving a monotone non-increasing curve +suitable for plotting optimisation progress. The x-axis is the fg-call index. + +For [`solve_ECG`](@ref) results this mirrors `convergence`. +""" function convergence_history(sr::SolverResults) return 1:length(sr.fg_history), sr.fg_history end +""" + correlation_function(sr; rmin=0.01, rmax=10.0, npoints=400, + coord_index=1, normalize=true) + +Compute the one-body radial density ``\\rho(r) = r^2 |\\psi_0(r)|^2`` along +a single Jacobi coordinate. + +# Arguments +- `sr` : [`SolverResults`](@ref) from either solver. +- `rmin`, `rmax` : radial grid range (a.u.). +- `npoints` : number of grid points. +- `coord_index` : which Jacobi coordinate to scan (default 1). +- `normalize` : if `true`, normalise so that ``\\int \\rho(r)\\,dr = 1``. + +Returns `(r_grid, ρ)` as plain `Vector{Float64}`. +""" function correlation_function( sr::SolverResults; rmin::Real = 0.01, diff --git a/src/variational.jl b/src/variational.jl index 2c32dfa..c6b86b0 100644 --- a/src/variational.jl +++ b/src/variational.jl @@ -123,9 +123,9 @@ Two loss types are available via the `loss_type` keyword: The `A` matrices of the basis functions are parameterised through their Cholesky factors (log-diagonal) which keeps them positive-definite for -any parameter vector, giving a smooth unconstrained problem. Shift -vectors are fixed at zero, which is appropriate for ground states with -spherical symmetry. +any parameter vector. Shift vectors `s` are also included in the +optimised parameter vector (unconstrained), so the full variational +freedom of each Gaussian is exploited. After optimisation the full generalised eigenvalue problem `Hc = λSc` is solved once to produce the ground-state energy and eigenvectors, so @@ -298,7 +298,13 @@ function solve_ECG_variational( verbose && @info "Starting variational ECG optimisation" n_basis = n n_params = length(θ0) loss_type - θ_opt, f_opt, _, _, _ = optimize(fg, θ0, method) + # OptimKit emits @warn for linesearch bisection failures that it handles + # gracefully internally. Suppress them to keep output clean. + θ_opt, f_opt, _, _, _ = Base.CoreLogging.with_logger( + Base.CoreLogging.ConsoleLogger(Base.stderr, Base.CoreLogging.Error) + ) do + optimize(fg, θ0, method) + end verbose && @info "Optimisation done" loss = f_opt diff --git a/test/test_variational.jl b/test/test_variational.jl index 4bb526d..75159bb 100644 --- a/test/test_variational.jl +++ b/test/test_variational.jl @@ -1,7 +1,7 @@ using Test using LinearAlgebra using FewBodyECG -import FewBodyECG: _jacobi_transform, _encode_basis, _decode_basis, _chol_to_params, _params_to_matrix +import FewBodyECG: _jacobi_transform, _encode_basis, _decode_basis, _chol_to_params, _params_to_matrix, convergence_history # --------------------------------------------------------------------------- # Shared 2-body (hydrogen) and 3-body (H⁻) operator fixtures @@ -239,3 +239,40 @@ end @test all(isfinite, rho) @test all(rho .>= 0.0) end + +# --------------------------------------------------------------------------- +# convergence_history +# --------------------------------------------------------------------------- + +@testset "convergence_history returns correct axes" begin + ops = _hydrogen_ops() + sr = solve_ECG_variational(ops, 5; + scale = 1.0, max_iterations = 30, verbose = false + ) + + xs, ys = convergence_history(sr) + @test length(xs) == length(ys) + @test length(xs) == length(sr.fg_history) + @test xs == 1:length(sr.fg_history) + @test ys === sr.fg_history + @test issorted(ys; rev = true) # monotone non-increasing by construction +end + +# --------------------------------------------------------------------------- +# Shift vectors are optimised (not pinned to zero) +# --------------------------------------------------------------------------- + +@testset "shift vectors are included in optimised parameters" begin + # _encode_basis should pack n_chol + n_dim params per Gaussian. + # For a 1-D (hydrogen) basis: n_chol=1, n_dim=1 → 2 params per function. + # The second param is the shift; _decode_basis should round-trip it. + ops = _hydrogen_ops() + sr = solve_ECG_variational(ops, 4; + scale = 1.0, max_iterations = 50, verbose = false + ) + # Each basis function has a 1-D shift vector stored in s. + for g in sr.basis_functions + @test length(g.s) == 1 + @test isfinite(g.s[1]) + end +end From 5e8b08731d7a6cf4c9381815ed07b6d95e360e18 Mon Sep 17 00:00:00 2001 From: MartinMikkelsen Date: Sat, 28 Feb 2026 17:39:21 +0100 Subject: [PATCH 4/8] updated examples --- Examples/HydrogenStates.jl | 144 +++++++++++++++++++++++ docs/Manifest.toml | 4 +- docs/Project.toml | 1 + docs/src/API.md | 1 + src/FewBodyECG.jl | 2 +- src/utils.jl | 12 +- src/variational.jl | 235 +++++++++++++++++++++++++++++++++++++ test/test_variational.jl | 94 ++++++++++++++- 8 files changed, 483 insertions(+), 10 deletions(-) create mode 100644 Examples/HydrogenStates.jl diff --git a/Examples/HydrogenStates.jl b/Examples/HydrogenStates.jl new file mode 100644 index 0000000..0cadb33 --- /dev/null +++ b/Examples/HydrogenStates.jl @@ -0,0 +1,144 @@ +using FewBodyECG +using LinearAlgebra +using Plots + +masses = [1.0e15, 1.0] +Λmat = Λ(masses) +_, U = _jacobi_transform(masses) +w = U' * Float64.([1, -1]) # electron-proton separation in Jacobi coords + +ops = Operator[KineticOperator(Λmat); CoulombOperator(-1.0, w)] + +sr_1s = solve_ECG_sequential(ops, 16; + n_candidates = 8, scale = 1.0, max_iterations_step = 120, verbose = true) + +E_1s = sr_1s.ground_state +println("\n E(1s) = $(round(E_1s; digits = 8)) Ha") +println(" E_exact(1s) = -0.50000000 Ha") +println(" |ΔE| = $(round(abs(E_1s + 0.5); sigdigits = 2))\n") + +a_p = [1.0] +s_zero = [0.0] + +# Log-spaced widths covering the spatial extent of the 2p orbital (peak ~4 a₀) +alphas_p = exp10.(range(log10(0.003), log10(3.0), length = 16)) +basis_p = GaussianBase[] +E_2p_conv = Float64[] +vecs_2p = Matrix{Float64}(undef, 0, 0) + +for (k, α) in enumerate(alphas_p) + push!(basis_p, Rank1Gaussian([α;;], a_p, s_zero)) + bset = BasisSet(basis_p) + H = build_hamiltonian_matrix(bset, ops) + S = build_overlap_matrix(bset) + vals, vecs = solve_generalized_eigenproblem(H, S) + E_k = minimum(vals) + push!(E_2p_conv, E_k) + global vecs_2p = vecs + println(" step $(lpad(k, 2)) α = $(rpad(round(α; digits = 4), 7))" * + " E = $(round(E_k; digits = 8))") +end + +E_2p = minimum(E_2p_conv) +c_2p = vecs_2p[:, 1] +println("\n E(2p) = $(round(E_2p; digits = 8)) Ha") +println(" E_exact(2p) = -0.12500000 Ha") +println(" |ΔE| = $(round(abs(E_2p + 0.125); sigdigits = 2))\n") + +a_d = [1.0] +# α range chosen for the 3d orbital spatial scale: peak of r⁶exp(−2αr²) is at +# r = √(3/α), so α ≈ 1/27 ≈ 0.037 for the 3d state (peak at r ≈ 9 a₀). +alphas_d = exp10.(range(log10(0.002), log10(0.8), length = 12)) +basis_d = GaussianBase[] +E_3d_conv = Float64[] +vecs_3d = Matrix{Float64}(undef, 0, 0) + +for (k, α) in enumerate(alphas_d) + push!(basis_d, Rank2Gaussian([α;;], a_d, a_d, s_zero)) + bset = BasisSet(basis_d) + H = build_hamiltonian_matrix(bset, ops) + S = build_overlap_matrix(bset) + vals, vecs = solve_generalized_eigenproblem(H, S) + E_k = minimum(vals) + push!(E_3d_conv, E_k) + global vecs_3d = vecs + note = E_k < -1 / 18 ? " ← below −1/18" : "" + println(" step $(lpad(k, 2)) α = $(rpad(round(α; digits = 4), 7))" * + " E = $(round(E_k; digits = 6))" * note) +end + +E_rank2 = minimum(E_3d_conv) +c_rank2 = vecs_3d[:, 1] +println("\n E(Rank2) = $(round(E_rank2; digits = 6)) Ha") +println(" E_exact(3d) = $(round(-1/18; digits = 6)) Ha") +println(" (energy lies below 3d due to L=0 mixing)\n") + +function ψ_rank1(rval, c, bfs) + r_vec = [rval] + return sum( + c[i] * dot(bfs[i].a, r_vec) * + exp(-dot(r_vec, parent(bfs[i].A) * r_vec) + dot(bfs[i].s, r_vec)) + for i in eachindex(bfs) + ) +end + +function ψ_rank2(rval, c, bfs) + r_vec = [rval] + return sum( + c[i] * dot(bfs[i].a, r_vec) * dot(bfs[i].b, r_vec) * + exp(-dot(r_vec, parent(bfs[i].A) * r_vec) + dot(bfs[i].s, r_vec)) + for i in eachindex(bfs) + ) +end + +function normalise(ρ, grid) + dr = step(grid) + return ρ ./ (sum(ρ) * dr) +end + +r_1s = range(0.01, 12.0, length = 600) +r_2p = range(0.01, 22.0, length = 600) +r_3d = range(0.01, 35.0, length = 600) + +ρ_ecg_1s = normalise([r^2 * abs2(ψ₀([r], sr_1s)) for r in r_1s], r_1s) +ρ_ecg_2p = normalise([r^2 * abs2(ψ_rank1(r, c_2p, basis_p)) for r in r_2p], r_2p) +ρ_ecg_rank2 = normalise([r^2 * abs2(ψ_rank2(r, c_rank2, basis_d)) for r in r_3d], r_3d) + +ρ_exact_1s = normalise([r^2 * exp(-2r) for r in r_1s], r_1s) +ρ_exact_2p = normalise([r^4 * exp(-r) for r in r_2p], r_2p) +ρ_exact_3d = normalise([r^6 * exp(-2r / 3) for r in r_3d], r_3d) + +p = plot( + layout = (3, 1), + size = (720, 1000), + left_margin = 6Plots.mm, + bottom_margin = 4Plots.mm, + legend = :topright, +) + +plot!(p[1], collect(r_1s), ρ_ecg_1s; + label = "ECG Rank0 (16 fn.)", lw = 2.5, color = :steelblue) +plot!(p[1], collect(r_1s), ρ_exact_1s; + label = "Exact 1s", lw = 1.8, ls = :dash, color = :black) +xlabel!(p[1], "r (a.u.)") +ylabel!(p[1], "r²|ψ(r)|²") +title!(p[1], "1s (L=0) E = $(round(E_1s; digits=6)) Ha | exact = −0.5") + +plot!(p[2], collect(r_2p), ρ_ecg_2p; + label = "ECG Rank1 (16 fn.)", lw = 2.5, color = :tomato) +plot!(p[2], collect(r_2p), ρ_exact_2p; + label = "Exact 2p", lw = 1.8, ls = :dash, color = :black) +xlabel!(p[2], "r (a.u.)") +ylabel!(p[2], "r²|ψ(r)|²") +title!(p[2], "2p (L=1) E = $(round(E_2p; digits=6)) Ha | exact = −0.125") + +plot!(p[3], collect(r_3d), ρ_ecg_rank2; + label = "ECG Rank2 (12 fn.)", lw = 2.5, color = :seagreen) +plot!(p[3], collect(r_3d), ρ_exact_3d; + label = "Exact 3d shape (r⁶ e^{−2r/3})", lw = 1.8, ls = :dash, color = :black) +xlabel!(p[3], "r (a.u.)") +ylabel!(p[3], "r²|ψ(r)|²") +title!(p[3], "d-wave-like Rank2 E = $(round(E_rank2; digits=5)) Ha |" * + " exact 3d = $(round(-1/18; digits=5)) (L=0/L=2 mixed)") + +display(p) diff --git a/docs/Manifest.toml b/docs/Manifest.toml index 9bc62dd..1a4f43b 100644 --- a/docs/Manifest.toml +++ b/docs/Manifest.toml @@ -2,7 +2,7 @@ julia_version = "1.12.5" manifest_format = "2.0" -project_hash = "3d39bcf94fc41feb6a82ef07b456752f08ca6ada" +project_hash = "c38ab8011052630e07e27b80b1686e1283e03d26" [[deps.ADTypes]] git-tree-sha1 = "f7304359109c768cf32dc5fa2d371565bb63b68a" @@ -382,7 +382,7 @@ uuid = "b22a6f82-2f65-5046-a5b2-351ab43fb4e5" version = "8.0.1+0" [[deps.FewBodyECG]] -deps = ["FewBodyHamiltonians", "ForwardDiff", "LinearAlgebra", "Logging", "OptimKit", "QuasiMonteCarlo", "SpecialFunctions"] +deps = ["FewBodyHamiltonians", "ForwardDiff", "LinearAlgebra", "OptimKit", "QuasiMonteCarlo", "SpecialFunctions"] path = ".." uuid = "083b1810-24a1-4a79-9a41-145bb2bb8ceb" version = "1.0.5" diff --git a/docs/Project.toml b/docs/Project.toml index 2e161fa..f56b1df 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -5,6 +5,7 @@ FewBodyHamiltonians = "3a126c26-e5d7-4a95-83c3-3b69f8a11ded" Optim = "429524aa-4258-5aef-a3af-852621145aeb" Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" QuasiMonteCarlo = "8a4e6c94-4038-4cdc-81c3-7e6ffdb2a71b" +SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" [compat] Optim = "1.7.8" diff --git a/docs/src/API.md b/docs/src/API.md index 2e9efab..7ac082f 100644 --- a/docs/src/API.md +++ b/docs/src/API.md @@ -5,6 +5,7 @@ ```@docs solve_ECG solve_ECG_variational +solve_ECG_sequential ``` ## Operators diff --git a/src/FewBodyECG.jl b/src/FewBodyECG.jl index 9a48840..0daaf5c 100644 --- a/src/FewBodyECG.jl +++ b/src/FewBodyECG.jl @@ -17,7 +17,7 @@ export build_hamiltonian_matrix, build_overlap_matrix, solve_generalized_eigenpr export ψ₀, SolverResults, convergence, convergence_history, correlation_function, ψ -export solve_ECG_variational +export solve_ECG_variational, solve_ECG_sequential include("types.jl") include("coordinates.jl") diff --git a/src/utils.jl b/src/utils.jl index 47267a8..a20a516 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -1,7 +1,7 @@ """ SolverResults -Output of [`solve_ECG`](@ref) and [`solve_ECG_variational`](@ref). +Output of [`solve_ECG`](@ref), [`solve_ECG_variational`](@ref), and [`solve_ECG_sequential`](@ref). # Fields | field | type | description | @@ -9,13 +9,13 @@ Output of [`solve_ECG`](@ref) and [`solve_ECG_variational`](@ref). | `basis_functions` | `Vector{GaussianBase}` | optimised basis | | `n_basis` | `Int` | number of accepted/optimised functions | | `operators` | `Vector{Operator}` | kinetic + Coulomb operators passed to the solver | -| `method` | `Symbol` | `:quasirandom`, `:random`, or `:variational` | -| `sampler` | `DeterministicSamplingAlgorithm` | QMC sampler used (placeholder for variational results) | +| `method` | `Symbol` | `:quasirandom`, `:random`, `:variational`, or `:sequential` | +| `sampler` | `DeterministicSamplingAlgorithm` | QMC sampler used (placeholder for variational/sequential results) | | `length_scale` | `Float64` | Gaussian width scale | | `ground_state` | `Float64` | lowest eigenvalue (ground-state energy in Hartree) | -| `energies` | `Vector{Float64}` | energy at each greedy step (stochastic) or `[ground_state]` (variational) | -| `eigenvectors` | `Vector{Matrix{Float64}}` | eigenvector matrices at each step | -| `fg_history` | `Vector{Float64}` | cumulative-minimum energy per objective call (variational) or mirrors `energies` (stochastic) | +| `energies` | `Vector{Float64}` | energy at each greedy/sequential step, or `[ground_state]` (variational) | +| `eigenvectors` | `Vector{Matrix{Float64}}` | eigenvector matrices (one per step for stochastic/sequential; one for variational) | +| `fg_history` | `Vector{Float64}` | cumulative-minimum energy per objective call (variational/sequential) or mirrors `energies` (stochastic) | Use [`convergence`](@ref), [`convergence_history`](@ref), [`correlation_function`](@ref), and [`ψ₀`](@ref) to analyse the result. diff --git a/src/variational.jl b/src/variational.jl index c6b86b0..3366898 100644 --- a/src/variational.jl +++ b/src/variational.jl @@ -334,3 +334,238 @@ function solve_ECG_variational( fg_history, ) end + +# --------------------------------------------------------------------------- +# Sequential (SVM-style) solver +# --------------------------------------------------------------------------- + +""" + solve_ECG_sequential(operators, n=50; kwargs...) -> SolverResults + +Build an ECG basis of `n` `Rank0Gaussian` functions using **sequential variational +optimisation** (the Stochastic Variational Method, SVM). + +At each step `k = 1, …, n`: +1. `n_candidates` quasi-random `Rank0Gaussian` functions are generated. +2. The candidate giving the lowest ground-state energy (evaluated without further + optimisation) is appended to the current basis. +3. **All** `k × n_per` parameters of the combined basis are jointly optimised + by an L-BFGS minimisation of the chosen `loss_type`. + +Compared to [`solve_ECG_variational`](@ref), which optimises all `n` functions +simultaneously from a cold start, the sequential approach: +* Avoids the high-dimensional landscape of a full cold start. +* Produces a monotone non-increasing convergence curve (`energies[k]` after each step). +* Closely mirrors the SVM algorithm of the ECG literature (see e.g. Suzuki & Varga 1998). + +# Arguments +- `operators` : `Vector{<:Operator}` — kinetic + Coulomb operators +- `n` : number of basis functions (default 50) + +# Keyword arguments +| keyword | default | description | +|:----------------------|:----------|:------------| +| `n_candidates` | `10` | candidates sampled per step; the one giving the lowest pre-optimisation energy is accepted | +| `loss_type` | `:energy` | `:energy` (λ_min) or `:trace` (Tr(S⁻¹H)) | +| `scale` | `0.2` | characteristic length scale for quasi-random Gaussian widths | +| `optimizer` | `nothing` | any OptimKit algorithm; `nothing` builds L-BFGS from `max_iterations_step` and `gradient_tol` | +| `max_iterations_step` | `100` | L-BFGS iterations per sequential step (ignored if `optimizer` given) | +| `gradient_tol` | `1e-6` | gradient-norm tolerance per step (ignored if `optimizer` given) | +| `regularization` | `1e-10` | Tikhonov shift added to S | +| `verbose` | `true` | print per-step info messages | + +# Example + +```julia +using FewBodyECG +masses = [1.0e15, 1.0, 1.0] # H⁻ (fixed nucleus + two electrons) +Λmat = Λ(masses) +_, U = _jacobi_transform(masses) +w_list = [[1,-1,0],[1,0,-1],[0,1,-1]] +w_raw = [U'*Float64.(w) for w in w_list] +ops = Operator[KineticOperator(Λmat); + [CoulombOperator(c,w) for (c,w) in zip([-1.,-1.,1.], w_raw)]...] + +sr = solve_ECG_sequential(ops, 30; scale=1.0, verbose=false) +println(sr.ground_state) # converges toward -0.52775... Ha +``` +""" +function solve_ECG_sequential( + operators::Vector{<:FewBodyHamiltonians.Operator}, + n::Int = 50; + n_candidates::Int = 10, + loss_type::Symbol = :energy, + scale::Real = 0.2, + optimizer = nothing, + max_iterations_step::Int = 100, + gradient_tol::Real = 1.0e-6, + regularization::Real = 1.0e-10, + verbose::Bool = true + ) + + loss_type in (:energy, :trace) || + throw(ArgumentError("loss_type must be :energy or :trace, got :$loss_type")) + + n_dim = size(first(op for op in operators if op isa KineticOperator).K, 1) + n_chol = n_dim * (n_dim + 1) ÷ 2 + n_per = n_chol + n_dim + w_list = [op.w for op in operators if op isa CoulombOperator] + + # Per-step optimiser: verbosity 0 — we emit our own step-level @info. + method = if optimizer !== nothing + optimizer + else + LBFGS(; maxiter = max_iterations_step, gradtol = float(gradient_tol), + verbosity = 0) + end + + energy_log = Float64[] # all fg values across all steps → cummin history + E_history = Float64[] # ground-state energy after each step's optimisation + θ_running = Float64[] # parameter vector; grows by n_per each step + + verbose && @info "Starting sequential ECG" n_basis = n n_candidates n_per_gaussian = n_per loss_type + + for step in 1:n + k = step # basis size after this step + + # ── candidate selection ────────────────────────────────────────────── + # Sample n_candidates quasi-random Gaussians; pick the one giving the + # lowest ground-state energy before optimisation. + best_E_cand = Inf + best_θ_cand = Float64[] + + for c in 1:n_candidates + attempt = (step - 1) * n_candidates + c + bij = generate_bij(:quasirandom, attempt, length(w_list), float(scale)) + A = _generate_A_matrix(bij, w_list) + s = generate_shift(:quasirandom, attempt, n_dim, float(scale)) + cand = Rank0Gaussian(A, s) + θ_c = _encode_basis(BasisSet([cand])) + θ_t = [θ_running; θ_c] + try + b_t = _decode_basis(θ_t, k, n_dim) + H_t = build_hamiltonian_matrix(b_t, operators) + S_t = build_overlap_matrix(b_t) + ev_t, _ = solve_generalized_eigenproblem(H_t, S_t; regularization) + E_t = minimum(ev_t) + if E_t < best_E_cand + best_E_cand = E_t + best_θ_cand = θ_c + end + catch + continue + end + end + + isempty(best_θ_cand) && + error("All $n_candidates candidates failed at sequential step $step") + + θ_running = [θ_running; best_θ_cand] + + # ── optimise all k-function parameters ────────────────────────────── + _chunk = min(n_per * 5, length(θ_running)) + _grad_cfg = ForwardDiff.GradientConfig(nothing, θ_running, + ForwardDiff.Chunk(_chunk)) + step_log = Float64[] + + # Build the fg closure for the current k-function basis. + # Capture k, _grad_cfg, step_log, and other constants by reference; + # each loop iteration creates a fresh set of these locals. + fg_k = if loss_type === :energy + (θ::AbstractVector) -> begin + local val::Float64, c::Vector{Float64} + try + b = _decode_basis(θ, k, n_dim) + H = build_hamiltonian_matrix(b, operators) + S = build_overlap_matrix(b) + ev, ev_vecs = solve_generalized_eigenproblem(H, S; regularization) + idx = argmin(ev) + val = ev[idx] + c = ev_vecs[:, idx] + catch + return Inf, zeros(Float64, length(θ)) + end + isfinite(val) || return Inf, zeros(Float64, length(θ)) + push!(step_log, val) + G = try + ForwardDiff.gradient(θ, _grad_cfg, Val(false)) do θ_ad + b_ad = _decode_basis(θ_ad, k, n_dim) + H_ad = build_hamiltonian_matrix(b_ad, operators) + S_ad = build_overlap_matrix(b_ad) + dot(c, H_ad * c) - val * dot(c, S_ad * c) + end + catch + zeros(Float64, length(θ)) + end + return val, G + end + else # :trace + (θ::AbstractVector) -> begin + local val_t::Float64 + try + b = _decode_basis(θ, k, n_dim) + H = build_hamiltonian_matrix(b, operators) + S = build_overlap_matrix(b) + v = tr((S + regularization * I) \ H) + val_t = isfinite(v) ? v : Inf + catch + return Inf, zeros(Float64, length(θ)) + end + isfinite(val_t) || return Inf, zeros(Float64, length(θ)) + push!(step_log, val_t) + G = try + ForwardDiff.gradient(θ, _grad_cfg, Val(false)) do θ_ad + b_ad = _decode_basis(θ_ad, k, n_dim) + H_ad = build_hamiltonian_matrix(b_ad, operators) + S_ad = build_overlap_matrix(b_ad) + tr((S_ad + regularization * I) \ H_ad) + end + catch + zeros(Float64, length(θ)) + end + return val_t, G + end + end + + θ_opt, _, _, _, _ = Base.CoreLogging.with_logger( + Base.CoreLogging.ConsoleLogger(Base.stderr, Base.CoreLogging.Error) + ) do + optimize(fg_k, θ_running, method) + end + θ_running = θ_opt + append!(energy_log, step_log) + + # Record the ground-state energy after this step's full optimisation. + b_k = _decode_basis(θ_running, k, n_dim) + H_k = build_hamiltonian_matrix(b_k, operators) + S_k = build_overlap_matrix(b_k) + ev_k, _ = solve_generalized_eigenproblem(H_k, S_k; regularization) + push!(E_history, minimum(ev_k)) + + verbose && @info "Step $step/$n" E₀ = last(E_history) fg_evals = length(step_log) + end + + # ── final reconstruction ───────────────────────────────────────────────── + basis_opt = _decode_basis(θ_running, n, n_dim) + H_opt = build_hamiltonian_matrix(basis_opt, operators) + S_opt = build_overlap_matrix(basis_opt) + evals, evecs = solve_generalized_eigenproblem(H_opt, S_opt) + ground_state = minimum(evals) + + verbose && @info "Sequential ECG complete" E₀ = ground_state n_basis = n + + fg_history = isempty(energy_log) ? Float64[] : accumulate(min, energy_log) + + return SolverResults( + Vector{GaussianBase}(basis_opt.functions), + n, + operators, + :sequential, + HaltonSample(), + float(scale), + ground_state, + E_history, # one energy per sequential step — use convergence() + [evecs], + fg_history, + ) +end diff --git a/test/test_variational.jl b/test/test_variational.jl index 75159bb..0c6bb0e 100644 --- a/test/test_variational.jl +++ b/test/test_variational.jl @@ -1,7 +1,7 @@ using Test using LinearAlgebra using FewBodyECG -import FewBodyECG: _jacobi_transform, _encode_basis, _decode_basis, _chol_to_params, _params_to_matrix, convergence_history +import FewBodyECG: _jacobi_transform, _encode_basis, _decode_basis, _chol_to_params, _params_to_matrix, convergence_history, solve_ECG_sequential # --------------------------------------------------------------------------- # Shared 2-body (hydrogen) and 3-body (H⁻) operator fixtures @@ -276,3 +276,95 @@ end @test isfinite(g.s[1]) end end + +# --------------------------------------------------------------------------- +# solve_ECG_sequential tests +# --------------------------------------------------------------------------- + +@testset "solve_ECG_sequential argument validation" begin + ops = _hydrogen_ops() + @test_throws ArgumentError solve_ECG_sequential( + ops, 3; loss_type = :bad, verbose = false + ) +end + +@testset "solve_ECG_sequential returns valid SolverResults" begin + ops = _hydrogen_ops() + sr = solve_ECG_sequential(ops, 4; + n_candidates = 3, scale = 1.0, max_iterations_step = 10, verbose = false + ) + + @test sr isa SolverResults + @test sr.n_basis == 4 + @test length(sr.basis_functions) == 4 + @test isfinite(sr.ground_state) + @test sr.ground_state < 0.0 + @test sr.method === :sequential + # energies has one entry per sequential step + @test length(sr.energies) == 4 + # eigenvectors: one final matrix + @test length(sr.eigenvectors) == 1 + @test size(sr.eigenvectors[1]) == (4, 4) + @test !isempty(sr.fg_history) +end + +@testset "solve_ECG_sequential convergence is monotone" begin + # By the variational principle, adding a linearly independent function + # and then re-optimising cannot raise the ground-state energy. + ops = _hydrogen_ops() + sr = solve_ECG_sequential(ops, 6; + n_candidates = 3, scale = 1.0, max_iterations_step = 20, verbose = false + ) + for i in 2:length(sr.energies) + @test sr.energies[i] <= sr.energies[i - 1] + 1.0e-8 + end +end + +@testset "solve_ECG_sequential respects variational bound (hydrogen)" begin + ops = _hydrogen_ops() + E_exact = -0.5 # hydrogen 1s ground state + + sr = solve_ECG_sequential(ops, 6; + n_candidates = 5, scale = 1.0, max_iterations_step = 50, verbose = false + ) + + @test sr.ground_state >= E_exact - 1.0e-6 # cannot go below exact + @test sr.ground_state < E_exact + 0.01 # should be close with 6 functions +end + +@testset "solve_ECG_sequential convergence_history is monotone" begin + ops = _hydrogen_ops() + sr = solve_ECG_sequential(ops, 4; + n_candidates = 3, scale = 1.0, max_iterations_step = 15, verbose = false + ) + xs, ys = convergence_history(sr) + @test length(xs) == length(ys) + @test issorted(ys; rev = true) +end + +@testset "solve_ECG_sequential beats stochastic (hydrogen, same n)" begin + ops = _hydrogen_ops() + + sr_stoch = solve_ECG(ops, 6; scale = 1.0, verbose = false) + sr_seq = solve_ECG_sequential(ops, 6; + n_candidates = 5, scale = 1.0, max_iterations_step = 50, verbose = false + ) + + @test sr_seq.ground_state <= sr_stoch.ground_state + 1.0e-4 +end + +@testset "ψ₀ and correlation_function work with sequential SolverResults" begin + ops = _hminus_ops() + sr = solve_ECG_sequential(ops, 4; + n_candidates = 3, scale = 1.0, max_iterations_step = 10, verbose = false + ) + + r_vec = [0.5, 0.3] + psi = ψ₀(r_vec, sr; state = 1) + @test isfinite(psi) + + r_grid, rho = correlation_function(sr; npoints = 30) + @test length(r_grid) == 30 + @test all(isfinite, rho) + @test all(rho .>= 0.0) +end From 858b6917126c084218eede064d229f241ff369ec Mon Sep 17 00:00:00 2001 From: MartinMikkelsen Date: Sat, 28 Feb 2026 20:37:23 +0100 Subject: [PATCH 5/8] updated examples --- Examples/HydrogenAnion.jl | 2 +- Examples/HydrogenStates.jl | 71 ++++++++----- Examples/Hydrogen_p-wave.jl | 4 +- Examples/Positronium.jl | 2 +- Project.toml | 8 +- docs/src/API.md | 2 + src/FewBodyECG.jl | 1 + src/hamiltonian.jl | 2 +- src/matrix_elements.jl | 188 ++++++++++++++++++++++++----------- src/types.jl | 115 ++++++++++++++++++--- src/utils.jl | 26 ----- test/test_hydrogen.jl | 17 ++-- test/test_matrix_elements.jl | 129 ++++++++++++++++++++++++ test/test_types.jl | 32 ++++++ 14 files changed, 458 insertions(+), 141 deletions(-) diff --git a/Examples/HydrogenAnion.jl b/Examples/HydrogenAnion.jl index 56127c4..fd5ec20 100644 --- a/Examples/HydrogenAnion.jl +++ b/Examples/HydrogenAnion.jl @@ -19,7 +19,7 @@ ops = Operator[ (CoulombOperator(c, w) for (c, w) in zip(coeffs, w_raw))... ] -result = solve_ECG(ops, 250, scale = 1.0) +result = solve_ECG(ops, 250, scale = 1.0, verbose=false) E = -0.527751016523 ΔE = abs(result.ground_state - E) diff --git a/Examples/HydrogenStates.jl b/Examples/HydrogenStates.jl index 0cadb33..22776a5 100644 --- a/Examples/HydrogenStates.jl +++ b/Examples/HydrogenStates.jl @@ -1,7 +1,10 @@ using FewBodyECG +import Antique using LinearAlgebra using Plots +atom = Antique.HydrogenAtom() + masses = [1.0e15, 1.0] Λmat = Λ(masses) _, U = _jacobi_transform(masses) @@ -14,8 +17,8 @@ sr_1s = solve_ECG_sequential(ops, 16; E_1s = sr_1s.ground_state println("\n E(1s) = $(round(E_1s; digits = 8)) Ha") -println(" E_exact(1s) = -0.50000000 Ha") -println(" |ΔE| = $(round(abs(E_1s + 0.5); sigdigits = 2))\n") +println(" E_exact(1s) = $(round(Antique.E(atom; n = 1); digits = 8)) Ha") +println(" |ΔE| = $(round(abs(E_1s - Antique.E(atom; n = 1)); sigdigits = 2))\n") a_p = [1.0] s_zero = [0.0] @@ -42,19 +45,19 @@ end E_2p = minimum(E_2p_conv) c_2p = vecs_2p[:, 1] println("\n E(2p) = $(round(E_2p; digits = 8)) Ha") -println(" E_exact(2p) = -0.12500000 Ha") -println(" |ΔE| = $(round(abs(E_2p + 0.125); sigdigits = 2))\n") +println(" E_exact(2p) = $(round(Antique.E(atom; n = 2); digits = 8)) Ha") +println(" |ΔE| = $(round(abs(E_2p - Antique.E(atom; n = 2)); sigdigits = 2))\n") -a_d = [1.0] -# α range chosen for the 3d orbital spatial scale: peak of r⁶exp(−2αr²) is at -# r = √(3/α), so α ≈ 1/27 ≈ 0.037 for the 3d state (peak at r ≈ 9 a₀). -alphas_d = exp10.(range(log10(0.002), log10(0.8), length = 12)) +# Orthogonal polarization vectors define a pure d-wave channel. +a_d = reshape([1.0, 0.0, 0.0], 1, 3) +b_d = reshape([0.0, 0.0, 1.0], 1, 3) +alphas_d = exp10.(range(log10(0.002), log10(0.8), length = 24)) basis_d = GaussianBase[] E_3d_conv = Float64[] vecs_3d = Matrix{Float64}(undef, 0, 0) for (k, α) in enumerate(alphas_d) - push!(basis_d, Rank2Gaussian([α;;], a_d, a_d, s_zero)) + push!(basis_d, Rank2Gaussian([α;;], a_d, b_d, s_zero)) bset = BasisSet(basis_d) H = build_hamiltonian_matrix(bset, ops) S = build_overlap_matrix(bset) @@ -62,16 +65,15 @@ for (k, α) in enumerate(alphas_d) E_k = minimum(vals) push!(E_3d_conv, E_k) global vecs_3d = vecs - note = E_k < -1 / 18 ? " ← below −1/18" : "" println(" step $(lpad(k, 2)) α = $(rpad(round(α; digits = 4), 7))" * - " E = $(round(E_k; digits = 6))" * note) + " E = $(round(E_k; digits = 8))") end E_rank2 = minimum(E_3d_conv) c_rank2 = vecs_3d[:, 1] -println("\n E(Rank2) = $(round(E_rank2; digits = 6)) Ha") -println(" E_exact(3d) = $(round(-1/18; digits = 6)) Ha") -println(" (energy lies below 3d due to L=0 mixing)\n") +println("\n E(3d, Rank2 pure d) = $(round(E_rank2; digits = 8)) Ha") +println(" E_exact(3d) = $(round(Antique.E(atom; n = 3); digits = 8)) Ha") +println(" |ΔE| = $(round(abs(E_rank2 - Antique.E(atom; n = 3)); sigdigits = 2))\n") function ψ_rank1(rval, c, bfs) r_vec = [rval] @@ -83,10 +85,16 @@ function ψ_rank1(rval, c, bfs) end function ψ_rank2(rval, c, bfs) + # Directional profile for the pure d-wave basis: r̂ = (x+z)/√2. + θ_d, φ_d = π / 4, 0.0 + r_cart = rval .* [sin(θ_d) * cos(φ_d), sin(θ_d) * sin(φ_d), cos(θ_d)] r_vec = [rval] return sum( - c[i] * dot(bfs[i].a, r_vec) * dot(bfs[i].b, r_vec) * - exp(-dot(r_vec, parent(bfs[i].A) * r_vec) + dot(bfs[i].s, r_vec)) + c[i] * ( + bfs[i].a isa AbstractMatrix ? + dot(vec(bfs[i].a), r_cart) * dot(vec(bfs[i].b), r_cart) : + dot(bfs[i].a, r_vec) * dot(bfs[i].b, r_vec) + ) * exp(-dot(r_vec, parent(bfs[i].A) * r_vec) + dot(bfs[i].s, r_vec)) for i in eachindex(bfs) ) end @@ -104,9 +112,21 @@ r_3d = range(0.01, 35.0, length = 600) ρ_ecg_2p = normalise([r^2 * abs2(ψ_rank1(r, c_2p, basis_p)) for r in r_2p], r_2p) ρ_ecg_rank2 = normalise([r^2 * abs2(ψ_rank2(r, c_rank2, basis_d)) for r in r_3d], r_3d) -ρ_exact_1s = normalise([r^2 * exp(-2r) for r in r_1s], r_1s) -ρ_exact_2p = normalise([r^4 * exp(-r) for r in r_2p], r_2p) -ρ_exact_3d = normalise([r^6 * exp(-2r / 3) for r in r_3d], r_3d) +θ_p, φ_p = 0.0, 0.0 +θ_d, φ_d = π / 4, 0.0 + +ρ_exact_1s = normalise( + [r^2 * abs2(Antique.ψ(atom, r, 0.0, 0.0; n = 1, l = 0, m = 0)) for r in r_1s], + r_1s, +) +ρ_exact_2p = normalise( + [r^2 * abs2(Antique.ψ(atom, r, θ_p, φ_p; n = 2, l = 1, m = 0)) for r in r_2p], + r_2p, +) +ρ_exact_3d = normalise( + [r^2 * abs2(Antique.ψ(atom, r, θ_d, φ_d; n = 3, l = 2, m = 1)) for r in r_3d], + r_3d, +) p = plot( layout = (3, 1), @@ -119,26 +139,25 @@ p = plot( plot!(p[1], collect(r_1s), ρ_ecg_1s; label = "ECG Rank0 (16 fn.)", lw = 2.5, color = :steelblue) plot!(p[1], collect(r_1s), ρ_exact_1s; - label = "Exact 1s", lw = 1.8, ls = :dash, color = :black) + label = "Antique 1s", lw = 1.8, ls = :dash, color = :black) xlabel!(p[1], "r (a.u.)") ylabel!(p[1], "r²|ψ(r)|²") -title!(p[1], "1s (L=0) E = $(round(E_1s; digits=6)) Ha | exact = −0.5") +title!(p[1], "1s (L=0) E = $(round(E_1s; digits=6)) Ha | exact = $(round(Antique.E(atom; n = 1); digits = 6))") plot!(p[2], collect(r_2p), ρ_ecg_2p; label = "ECG Rank1 (16 fn.)", lw = 2.5, color = :tomato) plot!(p[2], collect(r_2p), ρ_exact_2p; - label = "Exact 2p", lw = 1.8, ls = :dash, color = :black) + label = "Antique 2p (θ=0)", lw = 1.8, ls = :dash, color = :black) xlabel!(p[2], "r (a.u.)") ylabel!(p[2], "r²|ψ(r)|²") -title!(p[2], "2p (L=1) E = $(round(E_2p; digits=6)) Ha | exact = −0.125") +title!(p[2], "2p (L=1) E = $(round(E_2p; digits=6)) Ha | exact = $(round(Antique.E(atom; n = 2); digits = 6))") plot!(p[3], collect(r_3d), ρ_ecg_rank2; label = "ECG Rank2 (12 fn.)", lw = 2.5, color = :seagreen) plot!(p[3], collect(r_3d), ρ_exact_3d; - label = "Exact 3d shape (r⁶ e^{−2r/3})", lw = 1.8, ls = :dash, color = :black) + label = "Antique 3d (θ=π/4, m=1)", lw = 1.8, ls = :dash, color = :black) xlabel!(p[3], "r (a.u.)") ylabel!(p[3], "r²|ψ(r)|²") -title!(p[3], "d-wave-like Rank2 E = $(round(E_rank2; digits=5)) Ha |" * - " exact 3d = $(round(-1/18; digits=5)) (L=0/L=2 mixed)") +title!(p[3], "pure d-wave Rank2 E = $(round(E_rank2; digits=6)) Ha | exact 3d = $(round(Antique.E(atom; n = 3); digits = 6))") display(p) diff --git a/Examples/Hydrogen_p-wave.jl b/Examples/Hydrogen_p-wave.jl index cf3d078..3ba4069 100644 --- a/Examples/Hydrogen_p-wave.jl +++ b/Examples/Hydrogen_p-wave.jl @@ -2,7 +2,7 @@ using FewBodyECG using LinearAlgebra using Plots -masses = [1.0e15, 1.0] +masses = [1e12, 1.0] Λmat = Λ(masses) kin = KineticOperator(Λmat) @@ -21,8 +21,6 @@ ops = Operator[ a_vec = [1.0] s_zero = [0.0] -# Use a range of Gaussian widths that spans the spatial extent of the 2p orbital. -# The 2p state is more diffuse than 1s, so we need wider Gaussians (smaller α). alphas = [0.005, 0.01, 0.02, 0.05, 0.1, 0.2, 0.5, 1.0] basis_fns = GaussianBase[] diff --git a/Examples/Positronium.jl b/Examples/Positronium.jl index 145afcd..1d5ac79 100644 --- a/Examples/Positronium.jl +++ b/Examples/Positronium.jl @@ -18,5 +18,5 @@ coulomb_ops = [CoulombOperator(c, w) for (c, w) in zip(coeffs, w_raw)] ops = Operator[kin; coulomb_ops...] scale = default_scale(masses) -result = solve_ECG(ops, 300, sampler = SobolSample(); scale = scale) +result = solve_ECG(ops, 300, sampler = SobolSample(); scale = scale, verbose=false) println("E ≈ ", result.ground_state) diff --git a/Project.toml b/Project.toml index 92dc1fc..008f4b4 100644 --- a/Project.toml +++ b/Project.toml @@ -4,6 +4,7 @@ version = "1.0.5" authors = ["Shuhei Ohno", "Martin Mikkelsen"] [deps] +Antique = "be6e5d0e-34a5-4c8f-af83-e1b5389203d8" FewBodyHamiltonians = "3a126c26-e5d7-4a95-83c3-3b69f8a11ded" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" @@ -12,19 +13,22 @@ QuasiMonteCarlo = "8a4e6c94-4038-4cdc-81c3-7e6ffdb2a71b" SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" [compat] +Antique = "0.12.0" Aqua = "0.8.13" FewBodyHamiltonians = "0.0.2" ForwardDiff = "1.3.2" LinearAlgebra = "1.7.3" OptimKit = "0.4.2" QuasiMonteCarlo = "0.3.3" +Random = "1.11.0" SpecialFunctions = "2.5.0" Test = "1.11.0" -julia = "1.7" +julia = "1.8" [extras] Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Test", "Aqua"] +test = ["Test", "Aqua", "Random"] diff --git a/docs/src/API.md b/docs/src/API.md index 7ac082f..84a2407 100644 --- a/docs/src/API.md +++ b/docs/src/API.md @@ -32,5 +32,7 @@ correlation_function _jacobi_transform GaussianBase Rank0Gaussian +Rank1Gaussian +Rank2Gaussian BasisSet ``` diff --git a/src/FewBodyECG.jl b/src/FewBodyECG.jl index 0daaf5c..8f161ba 100644 --- a/src/FewBodyECG.jl +++ b/src/FewBodyECG.jl @@ -1,6 +1,7 @@ module FewBodyECG using LinearAlgebra +import Antique using FewBodyHamiltonians export Λ, _jacobi_transform diff --git a/src/hamiltonian.jl b/src/hamiltonian.jl index 6997d05..3659677 100644 --- a/src/hamiltonian.jl +++ b/src/hamiltonian.jl @@ -267,7 +267,7 @@ function solve_ECG( # Condition check on the overlap submatrix. cond_S = cond(Symmetric(S_k)) if cond_S > max_condition - @warn "Overlap poorly conditioned (κ=$cond_S) at step $ki, rejecting" + if verbose == true @warn "Overlap poorly conditioned (κ=$cond_S) at step $ki, rejecting" end n_rejected += 1 continue end diff --git a/src/matrix_elements.jl b/src/matrix_elements.jl index 1a2e454..7e317c3 100644 --- a/src/matrix_elements.jl +++ b/src/matrix_elements.jl @@ -27,12 +27,26 @@ function _compute_matrix_element(bra::Rank1Gaussian, ket::Rank1Gaussian) Rt = R * t I0 = exp(0.25 * t' * R * t) * (π^n / det(S))^(3 / 2) - cov = 0.5 * dot(bra.a, R * ket.a) - mean_term = 0.25 * dot(bra.a, Rt) * dot(ket.a, Rt) + cov = 0.5 * _polar_contract(bra.a, R, ket.a) + mean_term = 0.25 * _polar_project_dot( + bra.a, + Rt, + ket.a, + Rt, + ) return (cov + mean_term) * I0 end function _compute_matrix_element(bra::Rank2Gaussian, ket::Rank2Gaussian) + _check_polarization_compat(bra.a, bra.b) + _check_polarization_compat(ket.a, ket.b) + _check_polarization_compat(bra.a, ket.a) + if (any(!iszero, bra.s) || any(!iszero, ket.s)) && _pol_ncomp(bra.a) > 1 + throw(ArgumentError( + "Rank2 overlap with nonzero shifts currently requires single-component polarizations" + )) + end + A, B = bra.A, ket.A a, b = bra.s, ket.s S = A + B @@ -45,12 +59,12 @@ function _compute_matrix_element(bra::Rank2Gaussian, ket::Rank2Gaussian) μ = 0.5 * Rt - Xμ = dot(bra.a, μ) - Yμ = dot(bra.b, μ) - Zμ = dot(ket.a, μ) - Wμ = dot(ket.b, μ) + Xμ = sum(_polar_projection(bra.a, μ)) + Yμ = sum(_polar_projection(bra.b, μ)) + Zμ = sum(_polar_projection(ket.a, μ)) + Wμ = sum(_polar_projection(ket.b, μ)) - cov(v, w) = 0.5 * dot(v, R * w) + cov(v, w) = 0.5 * _polar_contract(v, R, w) XY = cov(bra.a, bra.b) XZ = cov(bra.a, ket.a) XW = cov(bra.a, ket.b) @@ -94,19 +108,22 @@ function _compute_matrix_element(bra::Rank1Gaussian, ket::Rank1Gaussian, op::Kin end A, B = bra.A, ket.A - a = vec(bra.a) - b = vec(ket.a) + a = bra.a + b = ket.a K = op.K R = inv(A + B) n = size(R, 1) M0 = (π^n / det(A + B))^(3 / 2) - M1 = 0.5 * dot(a, R * b) * M0 + M1 = 0.5 * _polar_contract(a, R, b) * M0 T1 = 6 * tr(B * K * A * R) * M1 - T2 = dot(a, K * b) * M0 - T3 = (dot(b, R * A * K * B * R * a) + dot(a, R * A * K * B * R * b)) * M0 - T4 = -dot(a, R * A * K * b) * M0 - T5 = -dot(b, R * B * K * a) * M0 + T2 = _polar_contract(a, K, b) * M0 + T3 = ( + _polar_contract(b, R * A * K * B * R, a) + + _polar_contract(a, R * A * K * B * R, b) + ) * M0 + T4 = -_polar_contract(a, R * A * K, b) * M0 + T5 = -_polar_contract(b, R * B * K, a) * M0 return T1 + T2 + T3 + T4 + T5 end @@ -118,60 +135,94 @@ function _compute_matrix_element(bra::Rank2Gaussian, ket::Rank2Gaussian, op::Kin A, B = bra.A, ket.A a, b, c, d = bra.a, bra.b, ket.a, ket.b + _check_polarization_compat(a, b) + _check_polarization_compat(c, d) + _check_polarization_compat(a, c) K = op.K R = inv(A + B) n = size(R, 1) M0 = (π^n / det(A + B))^(3 / 2) M2 = 0.25 * ( - dot(a, R * b) * dot(c, R * d) + - dot(a, R * c) * dot(b, R * d) + - dot(a, R * d) * dot(b, R * c) + _polar_contract(a, R, b) * + _polar_contract(c, R, d) + + _polar_contract(a, R, c) * + _polar_contract(b, R, d) + + _polar_contract(a, R, d) * + _polar_contract(b, R, c) ) * M0 T1 = 6 * tr(B * K * A * R) * M2 T2 = 0.5 * ( - dot(a, K * c) * dot(b, R * d) + - dot(a, K * d) * dot(b, R * c) + - dot(b, K * c) * dot(a, R * d) + - dot(b, K * d) * dot(a, R * c) + _polar_contract(a, K, c) * + _polar_contract(b, R, d) + + _polar_contract(a, K, d) * + _polar_contract(b, R, c) + + _polar_contract(b, K, c) * + _polar_contract(a, R, d) + + _polar_contract(b, K, d) * + _polar_contract(a, R, c) ) * M0 RAKBR = R * A * K * B * R T3 = 0.5 * ( - dot(a, RAKBR * b) * dot(c, R * d) + - dot(a, RAKBR * c) * dot(b, R * d) + - dot(a, RAKBR * d) * dot(b, R * c) + - dot(b, RAKBR * a) * dot(c, R * d) + - dot(b, RAKBR * c) * dot(a, R * d) + - dot(b, RAKBR * d) * dot(a, R * c) + - dot(c, RAKBR * a) * dot(b, R * d) + - dot(c, RAKBR * b) * dot(a, R * d) + - dot(c, RAKBR * d) * dot(a, R * b) + - dot(d, RAKBR * a) * dot(b, R * c) + - dot(d, RAKBR * b) * dot(a, R * c) + - dot(d, RAKBR * c) * dot(a, R * b) + _polar_contract(a, RAKBR, b) * + _polar_contract(c, R, d) + + _polar_contract(a, RAKBR, c) * + _polar_contract(b, R, d) + + _polar_contract(a, RAKBR, d) * + _polar_contract(b, R, c) + + _polar_contract(b, RAKBR, a) * + _polar_contract(c, R, d) + + _polar_contract(b, RAKBR, c) * + _polar_contract(a, R, d) + + _polar_contract(b, RAKBR, d) * + _polar_contract(a, R, c) + + _polar_contract(c, RAKBR, a) * + _polar_contract(b, R, d) + + _polar_contract(c, RAKBR, b) * + _polar_contract(a, R, d) + + _polar_contract(c, RAKBR, d) * + _polar_contract(a, R, b) + + _polar_contract(d, RAKBR, a) * + _polar_contract(b, R, c) + + _polar_contract(d, RAKBR, b) * + _polar_contract(a, R, c) + + _polar_contract(d, RAKBR, c) * + _polar_contract(a, R, b) ) * M0 RAK = R * A * K T4 = -0.5 * ( - dot(c, RAK * d) * dot(a, R * b) + - dot(d, RAK * c) * dot(a, R * b) + - dot(a, RAK * c) * dot(d, R * b) + - dot(a, RAK * d) * dot(c, R * b) + - dot(b, RAK * c) * dot(d, R * a) + - dot(b, RAK * d) * dot(a, R * c) + _polar_contract(c, RAK, d) * + _polar_contract(a, R, b) + + _polar_contract(d, RAK, c) * + _polar_contract(a, R, b) + + _polar_contract(a, RAK, c) * + _polar_contract(d, R, b) + + _polar_contract(a, RAK, d) * + _polar_contract(c, R, b) + + _polar_contract(b, RAK, c) * + _polar_contract(d, R, a) + + _polar_contract(b, RAK, d) * + _polar_contract(a, R, c) ) * M0 KBR = K * B * R T5 = -0.5 * ( - dot(a, KBR * c) * dot(d, R * b) + - dot(a, KBR * d) * dot(c, R * b) + - dot(a, KBR * b) * dot(c, R * d) + - dot(b, KBR * c) * dot(d, R * a) + - dot(b, KBR * d) * dot(c, R * a) + - dot(b, KBR * a) * dot(c, R * d) + _polar_contract(a, KBR, c) * + _polar_contract(d, R, b) + + _polar_contract(a, KBR, d) * + _polar_contract(c, R, b) + + _polar_contract(a, KBR, b) * + _polar_contract(c, R, d) + + _polar_contract(b, KBR, c) * + _polar_contract(d, R, a) + + _polar_contract(b, KBR, d) * + _polar_contract(c, R, a) + + _polar_contract(b, KBR, a) * + _polar_contract(c, R, d) ) * M0 return T1 + T2 + T3 + T4 + T5 @@ -204,10 +255,14 @@ function _compute_matrix_element(bra::Rank1Gaussian, ket::Rank1Gaussian, op::Cou n = size(R, 1) β = 1 / (dot(w, R * w)) M0 = (π^n / det(A + B))^(3 / 2) - M1 = 0.5 * dot(a, R * b) * M0 + M1 = 0.5 * _polar_contract(a, R, b) * M0 Rw = R * w - return op.coefficient * (2 * sqrt(β / π) * M1 - sqrt(β / π) * β / 3 * dot(a, Rw) * dot(Rw, b) * M0) + return op.coefficient * ( + 2 * sqrt(β / π) * M1 - + sqrt(β / π) * β / 3 * + _polar_project_dot(a, Rw, b, Rw) * M0 + ) end function _compute_matrix_element(bra::Rank2Gaussian, ket::Rank2Gaussian, op::CoulombOperator) @@ -217,6 +272,9 @@ function _compute_matrix_element(bra::Rank2Gaussian, ket::Rank2Gaussian, op::Cou A, B = bra.A, ket.A a, b, c, d = bra.a, bra.b, ket.a, ket.b + _check_polarization_compat(a, b) + _check_polarization_compat(c, d) + _check_polarization_compat(a, c) w = op.w R = inv(A + B) n = size(R, 1) @@ -226,27 +284,39 @@ function _compute_matrix_element(bra::Rank2Gaussian, ket::Rank2Gaussian, op::Cou Rw = R * w M2 = 0.25 * ( - dot(a, R * b) * dot(c, R * d) + - dot(a, R * c) * dot(b, R * d) + - dot(a, R * d) * dot(b, R * c) + _polar_contract(a, R, b) * + _polar_contract(c, R, d) + + _polar_contract(a, R, c) * + _polar_contract(b, R, d) + + _polar_contract(a, R, d) * + _polar_contract(b, R, c) ) * M0 term1 = 2 * sqrt(β / π) * M2 - q2_1 = dot(a, Rw) * dot(Rw, b) * dot(c, R * d) - q2_2 = dot(a, Rw) * dot(Rw, c) * dot(b, R * d) - q2_3 = dot(a, Rw) * dot(Rw, d) * dot(b, R * c) - q2_4 = dot(b, Rw) * dot(Rw, c) * dot(a, R * d) - q2_5 = dot(b, Rw) * dot(Rw, d) * dot(a, R * c) - q2_6 = dot(c, Rw) * dot(Rw, d) * dot(a, R * b) + q2_1 = _polar_project_dot(a, Rw, b, Rw) * + _polar_contract(c, R, d) + q2_2 = _polar_project_dot(a, Rw, c, Rw) * + _polar_contract(b, R, d) + q2_3 = _polar_project_dot(a, Rw, d, Rw) * + _polar_contract(b, R, c) + q2_4 = _polar_project_dot(b, Rw, c, Rw) * + _polar_contract(a, R, d) + q2_5 = _polar_project_dot(b, Rw, d, Rw) * + _polar_contract(a, R, c) + q2_6 = _polar_project_dot(c, Rw, d, Rw) * + _polar_contract(a, R, b) term2 = -2 * sqrt(β / π) * β / 3 * 0.25 * ( q2_1 + q2_2 + q2_3 + q2_4 + q2_5 + q2_6 ) * M0 - q4_1 = dot(a, Rw) * dot(Rw, b) * dot(c, Rw) * dot(Rw, d) - q4_2 = dot(a, Rw) * dot(Rw, c) * dot(b, Rw) * dot(Rw, d) - q4_3 = dot(a, Rw) * dot(Rw, d) * dot(b, Rw) * dot(Rw, c) + q4_1 = _polar_project_dot(a, Rw, b, Rw) * + _polar_project_dot(c, Rw, d, Rw) + q4_2 = _polar_project_dot(a, Rw, c, Rw) * + _polar_project_dot(b, Rw, d, Rw) + q4_3 = _polar_project_dot(a, Rw, d, Rw) * + _polar_project_dot(b, Rw, c, Rw) term3 = 2 * sqrt(β / π) * β^2 / 10 * 0.5 * ( q4_1 + q4_2 + q4_3 diff --git a/src/types.jl b/src/types.jl index 88004be..6765660 100644 --- a/src/types.jl +++ b/src/types.jl @@ -16,7 +16,7 @@ abstract type GaussianBase end Basis function ``g(\\mathbf{r}) = \\exp(-\\mathbf{r}^T A\\,\\mathbf{r} + \\mathbf{s}^T\\mathbf{r})``. # Fields -- `A` : symmetric positive-definite ``n_{\\text{dim}} \\times n_{\\text{dim}}`` matrix controlling the Gaussian width and inter-particle correlations. +- `A` : symmetric positive-definite ``n_{\\text{dim}} \\times n_{\\text{dim}}`` matrix controlling the Gaussian width and correlations. - `s` : shift vector ``\\mathbf{s} \\in \\mathbb{R}^{n_{\\text{dim}}}``; controls the location of the Gaussian maximum. """ struct Rank0Gaussian{T <: Real, M <: AbstractMatrix{T}, V <: AbstractVector{T}} <: GaussianBase @@ -29,28 +29,115 @@ struct Rank0Gaussian{T <: Real, M <: AbstractMatrix{T}, V <: AbstractVector{T}} end end -struct Rank1Gaussian{T <: Real, M <: AbstractMatrix{T}, V <: AbstractVector{T}} <: GaussianBase +const Polarization{T} = Union{AbstractVector{T}, AbstractMatrix{T}} + +_pol_nrows(a::AbstractVector) = length(a) +_pol_nrows(a::AbstractMatrix) = size(a, 1) +_pol_ncomp(a::AbstractVector) = 1 +_pol_ncomp(a::AbstractMatrix) = size(a, 2) +_pol_cols(a::AbstractVector) = reshape(a, :, 1) +_pol_cols(a::AbstractMatrix) = a + +_polarization_components(a::Union{AbstractVector, AbstractMatrix}) = _pol_ncomp(a) +_polarization_matrix(a::Union{AbstractVector, AbstractMatrix}) = _pol_cols(a) + +function _check_polarization_compat( + a::Union{AbstractVector, AbstractMatrix}, + b::Union{AbstractVector, AbstractMatrix} + ) + _pol_ncomp(a) == _pol_ncomp(b) || + throw(DimensionMismatch("polarizations must have the same number of components")) + return nothing +end + +function _polar_contract( + a::Union{AbstractVector, AbstractMatrix}, + M::AbstractMatrix, + b::Union{AbstractVector, AbstractMatrix} + ) + _check_polarization_compat(a, b) + A = _pol_cols(a) + B = _pol_cols(b) + return tr(transpose(A) * M * B) +end + +function _polar_projection( + a::Union{AbstractVector, AbstractMatrix}, + x::AbstractVector + ) + A = _pol_cols(a) + return transpose(A) * x +end + +function _polar_project_dot( + a::Union{AbstractVector, AbstractMatrix}, + x::AbstractVector, + b::Union{AbstractVector, AbstractMatrix}, + y::AbstractVector + ) + _check_polarization_compat(a, b) + return dot(_polar_projection(a, x), _polar_projection(b, y)) +end + +""" + Rank1Gaussian(A, a, s) + +Rank-1 (p-wave-like) ECG basis function with linear prefactor. + +`a` can be either: +- a vector of length `size(A,1)` (single polarization component), or +- a matrix of size `size(A,1) × ncomp` (multi-component polarization). +""" +struct Rank1Gaussian{ + T <: Real, + M <: AbstractMatrix{T}, + P <: Polarization{T}, + V <: AbstractVector{T}, +} <: GaussianBase A::Symmetric{T, M} - a::V + a::P s::V - function Rank1Gaussian(A::AbstractMatrix{T}, a::AbstractVector{T}, s::AbstractVector{T}) where {T <: Real} + function Rank1Gaussian(A::AbstractMatrix{T}, a::Polarization{T}, s::AbstractVector{T}) where {T <: Real} size(A, 1) == size(A, 2) || throw(ArgumentError("A must be square")) - (length(a) == size(A, 1) && length(s) == size(A, 1)) || - throw(ArgumentError("length(a) and length(s) must equal size(A,1)")) - return new{T, typeof(A), typeof(a)}(Symmetric(A), a, s) + _pol_nrows(a) == size(A, 1) || + throw(ArgumentError("size(a,1) (or length(a)) must equal size(A,1)")) + length(s) == size(A, 1) || + throw(ArgumentError("length(s) must equal size(A,1)")) + return new{T, typeof(A), typeof(a), typeof(s)}(Symmetric(A), a, s) end end -struct Rank2Gaussian{T <: Real, M <: AbstractMatrix{T}, V <: AbstractVector{T}} <: GaussianBase +""" + Rank2Gaussian(A, a, b, s) + +Rank-2 (d-wave-like) ECG basis function with quadratic prefactor. + +`a` and `b` can each be either vectors or matrices. Their first dimension must +match `size(A,1)`. For matrix polarizations, `a` and `b` must have the same +number of columns (`ncomp`), enabling multi-component pure d-wave channels. +""" +struct Rank2Gaussian{ + T <: Real, + M <: AbstractMatrix{T}, + P <: Polarization{T}, + Q <: Polarization{T}, + V <: AbstractVector{T}, +} <: GaussianBase A::Symmetric{T, M} - a::V - b::V + a::P + b::Q s::V - function Rank2Gaussian(A::AbstractMatrix{T}, a::AbstractVector{T}, b::AbstractVector{T}, s::AbstractVector{T}) where {T <: Real} + function Rank2Gaussian(A::AbstractMatrix{T}, a::Polarization{T}, b::Polarization{T}, s::AbstractVector{T}) where {T <: Real} size(A, 1) == size(A, 2) || throw(ArgumentError("A must be square")) - (length(a) == size(A, 1) && length(b) == size(A, 1) && length(s) == size(A, 1)) || - throw(ArgumentError("length(a), length(b), length(s) must equal size(A,1)")) - return new{T, typeof(A), typeof(a)}(Symmetric(A), a, b, s) + _pol_nrows(a) == size(A, 1) || + throw(ArgumentError("size(a,1) (or length(a)) must equal size(A,1)")) + _pol_nrows(b) == size(A, 1) || + throw(ArgumentError("size(b,1) (or length(b)) must equal size(A,1)")) + _pol_ncomp(a) == _pol_ncomp(b) || + throw(ArgumentError("a and b must have the same number of polarization components")) + length(s) == size(A, 1) || + throw(ArgumentError("length(s) must equal size(A,1)")) + return new{T, typeof(A), typeof(a), typeof(b), typeof(s)}(Symmetric(A), a, b, s) end end diff --git a/src/utils.jl b/src/utils.jl index a20a516..f295f2c 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -1,25 +1,3 @@ -""" - SolverResults - -Output of [`solve_ECG`](@ref), [`solve_ECG_variational`](@ref), and [`solve_ECG_sequential`](@ref). - -# Fields -| field | type | description | -|:------------------|:--------------------------------------|:------------| -| `basis_functions` | `Vector{GaussianBase}` | optimised basis | -| `n_basis` | `Int` | number of accepted/optimised functions | -| `operators` | `Vector{Operator}` | kinetic + Coulomb operators passed to the solver | -| `method` | `Symbol` | `:quasirandom`, `:random`, `:variational`, or `:sequential` | -| `sampler` | `DeterministicSamplingAlgorithm` | QMC sampler used (placeholder for variational/sequential results) | -| `length_scale` | `Float64` | Gaussian width scale | -| `ground_state` | `Float64` | lowest eigenvalue (ground-state energy in Hartree) | -| `energies` | `Vector{Float64}` | energy at each greedy/sequential step, or `[ground_state]` (variational) | -| `eigenvectors` | `Vector{Matrix{Float64}}` | eigenvector matrices (one per step for stochastic/sequential; one for variational) | -| `fg_history` | `Vector{Float64}` | cumulative-minimum energy per objective call (variational/sequential) or mirrors `energies` (stochastic) | - -Use [`convergence`](@ref), [`convergence_history`](@ref), [`correlation_function`](@ref), -and [`ψ₀`](@ref) to analyse the result. -""" struct SolverResults basis_functions::Vector{GaussianBase} n_basis::Int @@ -30,10 +8,6 @@ struct SolverResults ground_state::Float64 energies::Vector{Float64} eigenvectors::Vector{Matrix{Float64}} - # Per-call objective history (energy for :energy loss, trace for :trace loss). - # For solve_ECG this mirrors energies; for solve_ECG_variational it records - # the cumulative minimum at every primal fg evaluation so the curve is - # monotone and plottable via convergence_history(). fg_history::Vector{Float64} end diff --git a/test/test_hydrogen.jl b/test/test_hydrogen.jl index 913124b..6fff64a 100644 --- a/test/test_hydrogen.jl +++ b/test/test_hydrogen.jl @@ -66,22 +66,23 @@ end @test abs(E_best - E_exact) < 0.01 end -@testset "Rank2 Hamiltonian matrix is symmetric" begin - # Rank2 with scalar polarizations gives r²·exp(-αr²) (mixed s+d wave). - # Pure d-wave requires 3D polarization vectors (a·b=0), which is beyond scalar code. - # Here we just test that the Hamiltonian and overlap matrices are well-formed. - a_vec = [1.0] - b_vec = [1.0] - alphas = [0.1, 0.5, 1.0, 3.0] - basis_fns = [Rank2Gaussian([α;;], a_vec, b_vec, s_zero) for α in alphas] +@testset "Hydrogen d-wave (Rank2) pure-state convergence" begin + # Matrix polarizations with a⋅b = 0 enforce a pure d-wave channel. + a_mat = reshape([1.0, 0.0, 0.0], 1, 3) + b_mat = reshape([0.0, 1.0, 0.0], 1, 3) + alphas = exp10.(range(log10(0.002), log10(0.8), length = 24)) + basis_fns = [Rank2Gaussian([α;;], a_mat, b_mat, s_zero) for α in alphas] basis = BasisSet(basis_fns) H = build_hamiltonian_matrix(basis, ops) S = build_overlap_matrix(basis) + E_best = _solve_gep(H, S) + E_exact = -1 / 18 # 3d state @test H ≈ H' @test S ≈ S' @test all(diag(S) .> 0) + @test abs(E_best - E_exact) < 1.0e-4 end @testset "Rank1 overlap is correct" begin diff --git a/test/test_matrix_elements.jl b/test/test_matrix_elements.jl index 48a7101..0f8cc85 100644 --- a/test/test_matrix_elements.jl +++ b/test/test_matrix_elements.jl @@ -1,6 +1,7 @@ using Test using FewBodyECG using LinearAlgebra +using Random import FewBodyECG: _compute_matrix_element @@ -127,6 +128,134 @@ end @test isfinite(T) @test isfinite(Vval) end + + @testset "Matrix polarizations are backward compatible with single-column vectors" begin + A = [1.0 0.2; 0.2 1.5] + B = [0.9 0.1; 0.1 1.2] + s = [0.0, 0.0] + a = [0.5, -0.4] + b = [-0.2, 0.7] + c = [0.3, 0.6] + d = [-0.1, -0.8] + K = KineticOperator([0.5 0.0; 0.0 0.6]) + V = CoulombOperator(1.0, [1.0, 0.0]) + + g1v = Rank2Gaussian(A, a, b, s) + g2v = Rank2Gaussian(B, c, d, s) + g1m = Rank2Gaussian(A, reshape(a, :, 1), reshape(b, :, 1), s) + g2m = Rank2Gaussian(B, reshape(c, :, 1), reshape(d, :, 1), s) + + @test _compute_matrix_element(g1v, g2v) ≈ _compute_matrix_element(g1m, g2m) rtol = 1.0e-12 + @test _compute_matrix_element(g1v, g2v, K) ≈ _compute_matrix_element(g1m, g2m, K) rtol = 1.0e-12 + @test _compute_matrix_element(g1v, g2v, V) ≈ _compute_matrix_element(g1m, g2m, V) rtol = 1.0e-12 + end + + @testset "Incompatible polarization components throw DimensionMismatch" begin + A = [1.0 0.2; 0.2 1.5] + s = [0.0, 0.0] + + g1 = Rank1Gaussian(A, [0.5 0.1; -0.4 0.3], s) + g2 = Rank1Gaussian(A, [0.2; 0.7], s) + + @test_throws DimensionMismatch _compute_matrix_element(g1, g2) + @test_throws DimensionMismatch _compute_matrix_element(g1, g2, KineticOperator([0.5 0.0; 0.0 0.6])) + @test_throws DimensionMismatch _compute_matrix_element(g1, g2, CoulombOperator(1.0, [1.0, 0.0])) + end + + @testset "Rank2 kinetic/coulomb agree with paper formulas for matrix polarizations" begin + Random.seed!(7) + s(v, M, w) = tr(transpose(v) * M * w) + + function paper_rank2_kin(B, c, d, A, a, b, K) + R = inv(A + B) + n = size(R, 1) + M0 = (π^n / det(A + B))^(3 / 2) + M2 = 0.25 * (s(a, R, b) * s(c, R, d) + s(a, R, c) * s(b, R, d) + s(a, R, d) * s(b, R, c)) * M0 + T1 = 6 * tr(B * K * A * R) * M2 + T2 = 0.5 * (s(a, K, c) * s(b, R, d) + s(a, K, d) * s(b, R, c) + s(b, K, c) * s(a, R, d) + s(b, K, d) * s(a, R, c)) * M0 + + M = R * B * K * A * R + T3 = 0.5 * ( + s(a, M, b) * s(c, R, d) + s(a, M, c) * s(b, R, d) + s(a, M, d) * s(b, R, c) + + s(b, M, a) * s(c, R, d) + s(b, M, c) * s(a, R, d) + s(b, M, d) * s(a, R, c) + + s(c, M, a) * s(b, R, d) + s(c, M, b) * s(a, R, d) + s(c, M, d) * s(a, R, b) + + s(d, M, a) * s(b, R, c) + s(d, M, b) * s(a, R, c) + s(d, M, c) * s(a, R, b) + ) * M0 + + RBK = R * B * K + T4 = -0.5 * ( + s(a, RBK, b) * s(c, R, d) + s(b, RBK, a) * s(c, R, d) + + s(c, RBK, a) * s(b, R, d) + s(c, RBK, b) * s(a, R, d) + + s(d, RBK, a) * s(b, R, c) + s(d, RBK, b) * s(a, R, c) + ) * M0 + + KAR = K * A * R + T5 = -0.5 * ( + s(c, KAR, a) * s(b, R, d) + s(c, KAR, b) * s(a, R, d) + s(c, KAR, d) * s(a, R, b) + + s(d, KAR, a) * s(b, R, c) + s(d, KAR, b) * s(a, R, c) + s(d, KAR, c) * s(a, R, b) + ) * M0 + + return T1 + T2 + T3 + T4 + T5 + end + + function paper_rank2_coul(B, c, d, A, a, b, w, coef) + R = inv(A + B) + n = size(R, 1) + M0 = (π^n / det(A + B))^(3 / 2) + β = 1 / dot(w, R * w) + Rw = R * w + proj(x, y) = dot(transpose(x) * Rw, transpose(y) * Rw) + + M2 = 0.25 * (s(a, R, b) * s(c, R, d) + s(a, R, c) * s(b, R, d) + s(a, R, d) * s(b, R, c)) * M0 + term1 = 2 * sqrt(β / π) * M2 + + term2 = -2 * sqrt(β / π) * β / 3 * 0.25 * ( + proj(a, b) * s(c, R, d) + + proj(a, c) * s(b, R, d) + + proj(a, d) * s(b, R, c) + + proj(b, c) * s(a, R, d) + + proj(b, d) * s(a, R, c) + + proj(c, d) * s(a, R, b) + ) * M0 + + term3 = 2 * sqrt(β / π) * β^2 / 10 * 0.5 * ( + proj(a, b) * proj(c, d) + + proj(a, c) * proj(b, d) + + proj(a, d) * proj(b, c) + ) * M0 + + return coef * (term1 + term2 + term3) + end + + for n in (2, 3), _ in 1:3 + X = randn(n, n) + Y = randn(n, n) + Z = randn(n, n) + A = X' * X + I + B = Y' * Y + I + K = Z' * Z + w = randn(n) + + a = randn(n, 3) + b = randn(n, 3) + c = randn(n, 3) + d = randn(n, 3) + s0 = zeros(n) + + bra = Rank2Gaussian(A, a, b, s0) + ket = Rank2Gaussian(B, c, d, s0) + opK = KineticOperator(K) + opV = CoulombOperator(-1.3, w) + + gotK = _compute_matrix_element(bra, ket, opK) + gotV = _compute_matrix_element(bra, ket, opV) + refK = paper_rank2_kin(B, c, d, A, a, b, K) + refV = paper_rank2_coul(B, c, d, A, a, b, w, -1.3) + + @test gotK ≈ refK rtol = 1.0e-10 atol = 1.0e-12 + @test gotV ≈ refV rtol = 1.0e-10 atol = 1.0e-12 + end + end end @testset "Kinetic Energy ⟨g′|K|g⟩" begin diff --git a/test/test_types.jl b/test/test_types.jl index 2190e2d..71bf771 100644 --- a/test/test_types.jl +++ b/test/test_types.jl @@ -40,6 +40,20 @@ end end +@testset "Rank1Gaussian matrix polarization constructor" begin + A = [3.0 0.0; 0.0 4.0] + a_mat = [0.1 0.3 0.5; 0.2 0.4 0.6] + s = [1.0, 2.0] + + g1 = Rank1Gaussian(A, a_mat, s) + @test isa(g1, Rank1Gaussian) + @test g1.a == a_mat + @test g1.s == s + + @test_throws ArgumentError Rank1Gaussian(A, [0.1 0.3], s) + @test_throws ArgumentError Rank1Gaussian(A, a_mat, [1.0]) +end + @testset "Rank0Gaussian constructor and validate!" begin A = [2.0 0.0; 0.0 3.0] @@ -104,6 +118,24 @@ end @test_throws LinearAlgebra.PosDefException validate!(g2_indef) end +@testset "Rank2Gaussian matrix polarization constructor" begin + A = [4.0 0.0; 0.0 5.0] + a_mat = [0.1 0.3 0.5; 0.2 0.4 0.6] + b_mat = [0.7 0.9 1.1; 0.8 1.0 1.2] + s = [1.0, 2.0] + + g2 = Rank2Gaussian(A, a_mat, b_mat, s) + @test isa(g2, Rank2Gaussian) + @test g2.a == a_mat + @test g2.b == b_mat + @test g2.s == s + + @test_throws ArgumentError Rank2Gaussian(A, [0.1 0.3], b_mat, s) + @test_throws ArgumentError Rank2Gaussian(A, a_mat, [0.7 0.9], s) + @test_throws ArgumentError Rank2Gaussian(A, a_mat, [0.7; 0.8], s) + @test_throws ArgumentError Rank2Gaussian(A, a_mat, b_mat[:, 1:2], s) +end + @testset "validate! for Rank0 and Rank1 positive/negative-definite" begin A_pd = [2.0 0.0; 0.0 2.0] s = [0.0, 0.0] From 1eb6630743a27ef6df4e464c3f6f1c7811e5519f Mon Sep 17 00:00:00 2001 From: MartinMikkelsen Date: Mon, 2 Mar 2026 12:27:19 +0100 Subject: [PATCH 6/8] updated documentation --- docs/Manifest.toml | 10 ++- docs/Project.toml | 1 + docs/src/examples.md | 169 ++++++++++++++++++++++++++++++++++++++++++- src/types.jl | 11 +-- src/utils.jl | 18 +++++ 5 files changed, 198 insertions(+), 11 deletions(-) diff --git a/docs/Manifest.toml b/docs/Manifest.toml index 1a4f43b..f841d48 100644 --- a/docs/Manifest.toml +++ b/docs/Manifest.toml @@ -2,7 +2,7 @@ julia_version = "1.12.5" manifest_format = "2.0" -project_hash = "c38ab8011052630e07e27b80b1686e1283e03d26" +project_hash = "bd47568e186b9d1f91b1bb3bf1f3e2d19478f534" [[deps.ADTypes]] git-tree-sha1 = "f7304359109c768cf32dc5fa2d371565bb63b68a" @@ -73,6 +73,12 @@ git-tree-sha1 = "9876e1e164b144ca45e9e3198d0b689cadfed9ff" uuid = "66dad0bd-aa9a-41b7-9441-69ab47430ed8" version = "1.1.3" +[[deps.Antique]] +deps = ["SpecialFunctions"] +git-tree-sha1 = "bac153342749d4fe1c91918918e5a7a07f51d7e1" +uuid = "be6e5d0e-34a5-4c8f-af83-e1b5389203d8" +version = "0.12.0" + [[deps.ArgTools]] uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f" version = "1.1.2" @@ -382,7 +388,7 @@ uuid = "b22a6f82-2f65-5046-a5b2-351ab43fb4e5" version = "8.0.1+0" [[deps.FewBodyECG]] -deps = ["FewBodyHamiltonians", "ForwardDiff", "LinearAlgebra", "OptimKit", "QuasiMonteCarlo", "SpecialFunctions"] +deps = ["Antique", "FewBodyHamiltonians", "ForwardDiff", "LinearAlgebra", "OptimKit", "QuasiMonteCarlo", "SpecialFunctions"] path = ".." uuid = "083b1810-24a1-4a79-9a41-145bb2bb8ceb" version = "1.0.5" diff --git a/docs/Project.toml b/docs/Project.toml index f56b1df..4b9b989 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -1,4 +1,5 @@ [deps] +Antique = "be6e5d0e-34a5-4c8f-af83-e1b5389203d8" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" FewBodyECG = "083b1810-24a1-4a79-9a41-145bb2bb8ceb" FewBodyHamiltonians = "3a126c26-e5d7-4a95-83c3-3b69f8a11ded" diff --git a/docs/src/examples.md b/docs/src/examples.md index dc17b85..854ee93 100644 --- a/docs/src/examples.md +++ b/docs/src/examples.md @@ -73,7 +73,174 @@ plot(xs, ys; xlabel="fg evaluations", ylabel="Energy (Ha)", hline!([E_exact]; label="Exact", linestyle=:dot, color=:black) ``` -### Warm start from stochastic result +## Sequential variational solver (`solve_ECG_sequential`) + +`solve_ECG_sequential` builds the basis one function at a time. At each step +it samples `n_candidates` quasi-random Gaussians, picks the one that lowers the +energy the most, appends it to the current basis, and then jointly optimises +all accumulated parameters with L-BFGS. This combines the diversity of +stochastic sampling with gradient-based refinement at every step. + +Here we apply it to the hydrogen atom ground state (1s): + +```@example example_seq +using FewBodyECG +using LinearAlgebra +using Plots + +masses = [1.0e15, 1.0] # hydrogen: heavy nucleus + electron +Λmat = Λ(masses) +_, U = _jacobi_transform(masses) +w = U' * Float64.([1, -1]) # electron–nucleus separation + +ops = Operator[KineticOperator(Λmat); CoulombOperator(-1.0, w)] + +sr = solve_ECG_sequential(ops, 12; + n_candidates = 8, scale = 1.0, max_iterations_step = 80, verbose = false) + +println("E(1s) = ", round(sr.ground_state; digits = 8), " Ha") +println("Exact = -0.50000000 Ha") + +n_steps, E_steps = convergence(sr) +plot(n_steps, E_steps; + xlabel = "Basis size", ylabel = "Energy (Ha)", + label = "Sequential ECG", lw = 2, marker = :circle) +hline!([-0.5]; label = "Exact", ls = :dash, color = :black) +``` + +--- + +## Higher angular momentum: `Rank1Gaussian` (p-wave) and `Rank2Gaussian` (d-wave) + +Rank0 Gaussians are spherically symmetric. Non-zero angular momentum states +require polynomial prefactors. + +### p-wave with `Rank1Gaussian` + +`Rank1Gaussian(A, a, s)` adds a linear prefactor `(a⋅r)` selecting a spatial +direction. For the hydrogen 2p state (exact energy −1/8 Ha): + +```@example example_rank1 +using FewBodyECG +using LinearAlgebra +using Plots + +masses = [1.0e15, 1.0] +Λmat = Λ(masses) +_, U = _jacobi_transform(masses) +w = U' * Float64.([1, -1]) +ops = Operator[KineticOperator(Λmat); CoulombOperator(-1.0, w)] + +a_p = [1.0] # polarisation along the single Jacobi coordinate +s_zero = [0.0] + +alphas = exp10.(range(log10(0.003), log10(3.0), length = 12)) +basis = GaussianBase[] +E_conv = Float64[] + +for α in alphas + push!(basis, Rank1Gaussian([α;;], a_p, s_zero)) + bset = BasisSet(basis) + H = build_hamiltonian_matrix(bset, ops) + S = build_overlap_matrix(bset) + vals, _ = solve_generalized_eigenproblem(H, S) + push!(E_conv, minimum(vals)) +end + +println("E(2p) = ", round(minimum(E_conv); digits = 8), " Ha") +println("Exact = -0.12500000 Ha") + +plot(1:length(E_conv), E_conv; + xlabel = "Basis size", ylabel = "Energy (Ha)", + label = "ECG Rank1 (2p)", lw = 2, marker = :circle) +hline!([-0.125]; label = "Exact 2p", ls = :dash, color = :black) +``` + +### d-wave with `Rank2Gaussian` + +`Rank2Gaussian(A, a, b, s)` adds a quadratic prefactor `(a⋅r)(b⋅r)`. For a +pure d-wave channel the two polarisation vectors must be **orthogonal**. In a +1D Jacobi system (two-body) the three Cartesian directions are encoded as +columns of a `1 × 3` polarisation matrix: + +```@example example_rank2 +using FewBodyECG +using LinearAlgebra +using Plots + +masses = [1.0e15, 1.0] +Λmat = Λ(masses) +_, U = _jacobi_transform(masses) +w = U' * Float64.([1, -1]) +ops = Operator[KineticOperator(Λmat); CoulombOperator(-1.0, w)] + +# Orthogonal Cartesian directions → pure d-wave channel (a ⊥ b) +a_d = reshape([1.0, 0.0, 0.0], 1, 3) +b_d = reshape([0.0, 0.0, 1.0], 1, 3) +s_zero = [0.0] + +alphas = exp10.(range(log10(0.002), log10(0.8), length = 16)) +basis = GaussianBase[] +E_conv = Float64[] + +for α in alphas + push!(basis, Rank2Gaussian([α;;], a_d, b_d, s_zero)) + bset = BasisSet(basis) + H = build_hamiltonian_matrix(bset, ops) + S = build_overlap_matrix(bset) + vals, _ = solve_generalized_eigenproblem(H, S) + push!(E_conv, minimum(vals)) +end + +println("E(3d) = ", round(minimum(E_conv); digits = 8), " Ha") +println("Exact = ", round(-1/18; digits = 8), " Ha") + +plot(1:length(E_conv), E_conv; + xlabel = "Basis size", ylabel = "Energy (Ha)", + label = "ECG Rank2 (pure d-wave)", lw = 2, marker = :circle) +hline!([-1/18]; label = "Exact 3d", ls = :dash, color = :black) +``` + +--- + +## Muonic molecule tdμ + +The muonic three-body molecule tdμ (triton + deuteron + muon) is a nuclear-scale +system with masses three orders of magnitude larger than an atomic system. A +much smaller Gaussian width scale (~0.03 in nuclear units) is required; this can +be passed directly via the `scale` keyword. + +```@example example_tdmu +using FewBodyECG +using LinearAlgebra +using Plots + +masses = [5496.918, 3670.481, 206.7686] # t, d, μ in electron masses +Λmat = Λ(masses) +_, U = _jacobi_transform(masses) + +w_pairs = [[1, -1, 0], [1, 0, -1], [0, 1, -1]] +w_raw = [U' * Float64.(w) for w in w_pairs] +coeffs = [+1.0, -1.0, -1.0] # t-d repulsion; t-μ and d-μ attraction + +ops = Operator[ + KineticOperator(Λmat); + [CoulombOperator(c, w) for (c, w) in zip(coeffs, w_raw)]... +] + +result = solve_ECG(ops, 100; scale = 0.03, verbose = false) + +println("E(tdμ) = ", round(result.ground_state; digits = 5), " Ha") +println("SVM ref = -111.36444 Ha (Suzuki & Varga 1998, Table 8.1)") + +n, E = convergence(result) +plot(n, E; + xlabel = "Basis size", ylabel = "Energy (Ha)", + label = "ECG stochastic", lw = 2) +hline!([-111.36444]; label = "SVM reference", ls = :dash, color = :black) +``` + +## Warm start: stochastic basis + variational refinement When the stochastic solver has already found a reasonable basis, refining it with `loss_type = :trace` (minimise `Tr(S⁻¹H)`) often converges faster than a diff --git a/src/types.jl b/src/types.jl index 6765660..fd75b9d 100644 --- a/src/types.jl +++ b/src/types.jl @@ -31,16 +31,11 @@ end const Polarization{T} = Union{AbstractVector{T}, AbstractMatrix{T}} -_pol_nrows(a::AbstractVector) = length(a) -_pol_nrows(a::AbstractMatrix) = size(a, 1) _pol_ncomp(a::AbstractVector) = 1 _pol_ncomp(a::AbstractMatrix) = size(a, 2) _pol_cols(a::AbstractVector) = reshape(a, :, 1) _pol_cols(a::AbstractMatrix) = a -_polarization_components(a::Union{AbstractVector, AbstractMatrix}) = _pol_ncomp(a) -_polarization_matrix(a::Union{AbstractVector, AbstractMatrix}) = _pol_cols(a) - function _check_polarization_compat( a::Union{AbstractVector, AbstractMatrix}, b::Union{AbstractVector, AbstractMatrix} @@ -99,7 +94,7 @@ struct Rank1Gaussian{ s::V function Rank1Gaussian(A::AbstractMatrix{T}, a::Polarization{T}, s::AbstractVector{T}) where {T <: Real} size(A, 1) == size(A, 2) || throw(ArgumentError("A must be square")) - _pol_nrows(a) == size(A, 1) || + size(a, 1) == size(A, 1) || throw(ArgumentError("size(a,1) (or length(a)) must equal size(A,1)")) length(s) == size(A, 1) || throw(ArgumentError("length(s) must equal size(A,1)")) @@ -129,9 +124,9 @@ struct Rank2Gaussian{ s::V function Rank2Gaussian(A::AbstractMatrix{T}, a::Polarization{T}, b::Polarization{T}, s::AbstractVector{T}) where {T <: Real} size(A, 1) == size(A, 2) || throw(ArgumentError("A must be square")) - _pol_nrows(a) == size(A, 1) || + size(a, 1) == size(A, 1) || throw(ArgumentError("size(a,1) (or length(a)) must equal size(A,1)")) - _pol_nrows(b) == size(A, 1) || + size(b, 1) == size(A, 1) || throw(ArgumentError("size(b,1) (or length(b)) must equal size(A,1)")) _pol_ncomp(a) == _pol_ncomp(b) || throw(ArgumentError("a and b must have the same number of polarization components")) diff --git a/src/utils.jl b/src/utils.jl index f295f2c..fe03c40 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -1,3 +1,21 @@ +""" + SolverResults + +Container returned by all ECG solvers ([`solve_ECG`](@ref), +[`solve_ECG_variational`](@ref), [`solve_ECG_sequential`](@ref)). + +# Fields +- `basis_functions` : accepted basis functions. +- `n_basis` : number of accepted basis functions. +- `operators` : the operator list used during the solve. +- `method` : solver symbol (`:quasirandom`, `:random`, `:variational`, `:sequential`). +- `sampler` : quasi-/pseudo-random sampler used for basis generation. +- `length_scale` : Gaussian width scale passed at construction. +- `ground_state` : lowest eigenvalue (ground-state energy in a.u.). +- `energies` : energy after each accepted basis function (stochastic) or after each step (sequential). +- `eigenvectors` : list of eigenvector matrices; `eigenvectors[end][:, 1]` is the ground-state coefficient vector. +- `fg_history` : monotone-decreasing objective value after each gradient evaluation (variational solvers). +""" struct SolverResults basis_functions::Vector{GaussianBase} n_basis::Int From d49e5658d8e81ad9a3179a9bb835ffa758b2873f Mon Sep 17 00:00:00 2001 From: MartinMikkelsen Date: Fri, 6 Mar 2026 10:49:46 +0100 Subject: [PATCH 7/8] updated examples, added excited states --- Examples/Helium.jl | 66 +++++++++++ Examples/HydrogenAnion.jl | 37 ++++--- Examples/Hydrogen_p-wave.jl | 18 +-- Examples/Positronium.jl | 30 +++-- Examples/Variational.jl | 25 +---- "Examples/td\316\274.jl" | 17 +-- src/FewBodyECG.jl | 2 + src/hamiltonian.jl | 212 ++++++++++++++++++++++++++++++++++-- src/utils.jl | 8 +- src/variational.jl | 2 + test/test_utils.jl | 7 +- 11 files changed, 328 insertions(+), 96 deletions(-) create mode 100644 Examples/Helium.jl diff --git a/Examples/Helium.jl b/Examples/Helium.jl new file mode 100644 index 0000000..67f33c1 --- /dev/null +++ b/Examples/Helium.jl @@ -0,0 +1,66 @@ +using FewBodyECG +using Antique +using Plots +using QuasiMonteCarlo + +masses = [1e15, 1.0, 1.0] + +os = Operators(masses, [+2, -1, -1]) # nucleus (Z=2), e₁, e₂ +os += "Kinetic" +os += "Coulomb" # auto: nucleus-e₁ (-2), nucleus-e₂ (-2), e₁-e₂ (+1) + +E_gs_exact = -2.9037242 # 1s² ¹S (ground state) +E_ex_exact = -2.17523 # 1s2s ¹S (first excited singlet) + +println("Ground state (1s²)...") +result_gs = solve_ECG(os, 250; scale = 1.0, verbose = false) +ΔE_gs = result_gs.ground_state - E_gs_exact +println(" E = $(round(result_gs.ground_state; digits=6)) (exact: $E_gs_exact, error: $(round(ΔE_gs; digits=6)))") + +println("\nFirst excited ¹S state (1s2s)...") +result_ex = solve_ECG(os, 200; scale = 1.0, verbose = false, state = 2) +ΔE_ex = result_ex.ground_state - E_ex_exact +println(" E = $(round(result_ex.ground_state; digits=6)) (exact: $E_ex_exact, error: $(round(ΔE_ex; digits=6)))") + +n_gs, E_gs = convergence(result_gs) +n_ex, E_ex = convergence(result_ex) + +p1 = plot(n_gs, E_gs, + label = "1s² (ground)", lw = 2, + xlabel = "Basis size", ylabel = "E (Ha)", + title = "Helium convergence") +plot!(p1, n_ex, E_ex, label = "1s2s (excited)", lw = 2, ls = :dash) +hline!(p1, [E_gs_exact, E_ex_exact], ls = :dot, color = :gray, label = "Exact") +display(p1) + +r_gs, ρ_gs = correlation_function(result_gs; rmax = 5.0) +r_ex, ρ_ex = correlation_function(result_ex; rmax = 5.0) + +p2 = plot(r_gs, ρ_gs, + label = "1s² (ground)", lw = 2, + xlabel = "r (a.u.)", ylabel = "r²|ψ(r)|²", + title = "Helium radial correlation") + plot!(p2, r_ex, ρ_ex, label = "1s2s (excited)", lw = 2, ls = :dash) +display(p2) + +HeP = HydrogenAtom(Z = 2) # atomic units: Eₕ=1, a₀=1, mₑ=1, ℏ=1 (defaults) +E_hep_exact = Antique.E(HeP; n = 1) # = -2.0 Ha + +os_hep = Operators([1e15, 1.0], [+2, -1]) +os_hep += "Kinetic" +os_hep += "Coulomb" + +result_hep = solve_ECG(os_hep, 250; scale = 0.5, verbose = false) +ΔE_hep = result_hep.ground_state - E_hep_exact +println(" E (ECG) = $(round(result_hep.ground_state; digits=8))") +println(" E (Antique)= $(round(E_hep_exact; digits=8)) error: $(round(ΔE_hep; sigdigits=3))") + +r_hep, ρ_ecg = correlation_function(result_hep; rmax = 3.0, npoints = 300) +ρ_antique = [r^2 * Antique.R(HeP, r; n = 1, l = 0)^2 for r in r_hep] + +p3 = plot(r_hep, ρ_ecg, + label = "ECG", lw = 2, + xlabel = "r (a.u.)", ylabel = "r²|ψ(r)|²", + title = "He⁺ 1s radial density: ECG vs Antique.jl") +plot!(p3, r_hep, ρ_antique, label = "Antique (exact)", lw = 2, ls = :dash) +display(p3) diff --git a/Examples/HydrogenAnion.jl b/Examples/HydrogenAnion.jl index fd5ec20..2ca227c 100644 --- a/Examples/HydrogenAnion.jl +++ b/Examples/HydrogenAnion.jl @@ -5,25 +5,26 @@ using QuasiMonteCarlo masses = [1.0e15, 1.0, 1.0] -Λmat = Λ(masses) -kin = KineticOperator(Λmat) -J, U = _jacobi_transform(masses) +os = Operators(masses, [+1, -1, -1]) # proton, e₁, e₂ +os += "Kinetic" +os += "Coulomb" -w_list = [[1, -1, 0], [1, 0, -1], [0, 1, -1]] +result = solve_ECG(os, 250, scale = 1.0, verbose=false) -w_raw = [U' * w for w in w_list] -coeffs = [-1.0, -1.0, +1.0] - -ops = Operator[ - kin; - (CoulombOperator(c, w) for (c, w) in zip(coeffs, w_raw))... -] - -result = solve_ECG(ops, 250, scale = 1.0, verbose=false) - -E = -0.527751016523 -ΔE = abs(result.ground_state - E) +E_exact = -0.527751016523 +ΔE = abs(result.ground_state - E_exact) @info "Energy difference" ΔE -n, E = convergence(result) -plot(n, E) +n_conv, E_conv = convergence(result) +p1 = plot(n_conv, E_conv, + xlabel = "Basis size", ylabel = "E (Ha)", + label = "Ground state", lw = 2, + title = "Hydrogen anion convergence") +display(p1) + +r_grid, ρ = correlation_function(result; rmax = 10.0) +p2 = plot(r_grid, ρ, + xlabel = "r (a.u.)", ylabel = "r²|ψ(r)|²", + label = "Hydrogen anion", lw = 2, + title = "Hydrogen radial correlation") +display(p2) diff --git a/Examples/Hydrogen_p-wave.jl b/Examples/Hydrogen_p-wave.jl index 3ba4069..30d275c 100644 --- a/Examples/Hydrogen_p-wave.jl +++ b/Examples/Hydrogen_p-wave.jl @@ -4,20 +4,10 @@ using Plots masses = [1e12, 1.0] -Λmat = Λ(masses) -kin = KineticOperator(Λmat) +os = Operators(masses) +os += "Kinetic" +os += "Coulomb", 1, 2, -1.0 # p-e (attraction) -J, U = _jacobi_transform(masses) - -w_raw = [U' * [1, -1]] -coeffs = [-1.0] - -ops = Operator[ - kin; - (CoulombOperator(c, w) for (c, w) in zip(coeffs, w_raw))... -] - -# Polarization vector for p-wave (selects one spatial direction) a_vec = [1.0] s_zero = [0.0] @@ -31,7 +21,7 @@ for (i, α) in enumerate(alphas) basis = BasisSet(basis_fns) - H = build_hamiltonian_matrix(basis, ops) + H = build_hamiltonian_matrix(basis, os) S = build_overlap_matrix(basis) global vals, vecs = solve_generalized_eigenproblem(H, S) diff --git a/Examples/Positronium.jl b/Examples/Positronium.jl index 1d5ac79..aa368a6 100644 --- a/Examples/Positronium.jl +++ b/Examples/Positronium.jl @@ -5,18 +5,26 @@ import FewBodyECG: default_scale, convergence masses = [1.0, 1.0, 1.0] -Λmat = Λ(masses) -kin = KineticOperator(Λmat) +os = Operators(masses, [+1, -1, -1]) # e⁺, e⁻, e⁻ +os += "Kinetic" +os += "Coulomb" -J, U = _jacobi_transform(masses) -w_list = [[1, -1, 0], [1, 0, -1], [0, 1, -1]] -w_raw = [U' * w for w in w_list] - -coeffs = [-1.0, -1.0, +1.0] -coulomb_ops = [CoulombOperator(c, w) for (c, w) in zip(coeffs, w_raw)] - -ops = Operator[kin; coulomb_ops...] scale = default_scale(masses) -result = solve_ECG(ops, 300, sampler = SobolSample(); scale = scale, verbose=false) +# Ps⁻ has only one bound state; for excited-state examples see Helium.jl. +result = solve_ECG(os, 300, sampler = SobolSample(); scale = scale, verbose=false, state = 1) println("E ≈ ", result.ground_state) + +n_conv, E_conv = convergence(result) +p1 = plot(n_conv, E_conv, + xlabel = "Basis size", ylabel = "E (Ha)", + label = "Ground state", lw = 2, + title = "Positronium convergence") +display(p1) + +r_grid, ρ = correlation_function(result; rmax = 15.0) +p2 = plot(r_grid, ρ, + xlabel = "r (a.u.)", ylabel = "r²|ψ(r)|²", + label = "Positronium", lw = 2, + title = "Positronium radial correlation") +display(p2) diff --git a/Examples/Variational.jl b/Examples/Variational.jl index fb8c7f7..1ac3306 100644 --- a/Examples/Variational.jl +++ b/Examples/Variational.jl @@ -24,14 +24,6 @@ using LinearAlgebra using QuasiMonteCarlo import FewBodyECG: default_scale, BasisSet, Rank0Gaussian -# ============================================================ -# 1. Hydrogen anion H⁻ -# ============================================================ - -println("=" ^ 60) -println("Hydrogen anion H⁻") -println("=" ^ 60) - masses_Hm = [1.0e15, 1.0, 1.0] # fixed nucleus + 2 electrons Λ_Hm = Λ(masses_Hm) _, U_Hm = _jacobi_transform(masses_Hm) @@ -48,19 +40,16 @@ ops_Hm = Operator[ E_exact_Hm = -0.527751016523 n = 30 -# --- (a) fresh optimisation with :energy loss --- println("\n(a) Fresh start, loss_type = :energy") sr_var = solve_ECG_variational(ops_Hm, n; scale = 1.0, max_iterations = 500, verbose = false) ΔE_var = sr_var.ground_state - E_exact_Hm println(" Variational E₀ = $(round(sr_var.ground_state, digits=8)) ΔE = $(round(ΔE_var, sigdigits=3))") -# --- stochastic baseline --- sr_stoch = solve_ECG(ops_Hm, n; scale = 1.0, verbose = false) ΔE_stoch = sr_stoch.ground_state - E_exact_Hm println(" Stochastic E₀ = $(round(sr_stoch.ground_state, digits=8)) ΔE = $(round(ΔE_stoch, sigdigits=3))") println(" Exact E₀ = $E_exact_Hm") -# --- (b) warm-start from stochastic result with :trace loss --- println("\n(b) Warm-start from stochastic, loss_type = :trace") basis0_Hm = BasisSet(Rank0Gaussian[sr_stoch.basis_functions...]) sr_warm = solve_ECG_variational(ops_Hm, n; @@ -74,19 +63,9 @@ println(" Warm-start E₀ = $(round(sr_warm.ground_state, digits=8)) ΔE = $ println(" Stochastic E₀ = $(round(sr_stoch.ground_state, digits=8)) ΔE = $(round(ΔE_stoch, sigdigits=3))") println(" Exact E₀ = $E_exact_Hm") -# --- downstream utilities work unchanged --- r_grid, ρ = correlation_function(sr_var; rmin = 0.01, rmax = 15.0, npoints = 200) println("\n Correlation function computed: $(length(r_grid)) points, max ρ at r = $(round(r_grid[argmax(ρ)], digits=3)) a.u.") -# ============================================================ -# 2. Muonic molecule tdμ -# ============================================================ - -println() -println("=" ^ 60) -println("Muonic molecule tdμ (triton + deuteron + muon)") -println("=" ^ 60) - masses_tdμ = [5496.918, 3670.481, 206.7686] # t, d, μ in electron masses Λ_tdμ = Λ(masses_tdμ) _, U_tdμ = _jacobi_transform(masses_tdμ) @@ -100,7 +79,7 @@ ops_tdμ = Operator[ ] E_exact_tdμ = -111.36444 -scale_tdμ = 0.03 # nuclear scale (much smaller than atomic) +scale_tdμ = 0.03 # nuclear scale n_tdμ = 25 println("\n(a) Fresh start, loss_type = :energy") @@ -119,7 +98,6 @@ println(" Exact E₀ = $E_exact_tdμ") using Plots import FewBodyECG: convergence_history -# --- Variational convergence (energy vs fg evaluations) --- n_fg, E_fg = convergence_history(sr_var) p1 = plot(n_fg, E_fg; label = "Variational", @@ -130,7 +108,6 @@ p1 = plot(n_fg, E_fg; ) hline!(p1, [E_exact_Hm]; label = "Exact", linestyle = :dot, color = :black, lw = 1) -# --- Stochastic greedy convergence (energy vs basis size) --- n_s, E_s = convergence(sr_stoch) p2 = plot(n_s, E_s; label = "Stochastic greedy", diff --git "a/Examples/td\316\274.jl" "b/Examples/td\316\274.jl" index 0b0d814..5e33602 100644 --- "a/Examples/td\316\274.jl" +++ "b/Examples/td\316\274.jl" @@ -7,18 +7,11 @@ import FewBodyECG: _generate_A_matrix masses = [5496.918, 3670.481, 206.7686] -Λmat = Λ(masses) -kin = KineticOperator(Λmat) +os = Operators(masses, [+1, +1, -1]) # triton, deuteron, muon +os += "Kinetic" +os += "Coulomb" # auto: t-d (+1), t-μ (-1), d-μ (-1) -J, U = _jacobi_transform(masses) - -w_pairs = [[1, -1, 0], [1, 0, -1], [0, 1, -1]] - -w_jac = [U' * Float64.(w) for w in w_pairs] - -coeffs = [+1.0, -1.0, -1.0] - -ops = Operator[kin; [CoulombOperator(c, w) for (c, w) in zip(coeffs, w_jac)]...] +w_jac = coulomb_weights(os) E_exact = -111.36444 @@ -46,7 +39,7 @@ function compute_energy(points, n_basis, b₀) end basis = BasisSet(basis_fns) - H = build_hamiltonian_matrix(basis, ops) + H = build_hamiltonian_matrix(basis, os) S = build_overlap_matrix(basis) try diff --git a/src/FewBodyECG.jl b/src/FewBodyECG.jl index 8f161ba..7d3287c 100644 --- a/src/FewBodyECG.jl +++ b/src/FewBodyECG.jl @@ -20,6 +20,8 @@ export ψ₀, SolverResults, convergence, convergence_history, correlation_funct export solve_ECG_variational, solve_ECG_sequential +export Operators, coulomb_weights + include("types.jl") include("coordinates.jl") include("matrix_elements.jl") diff --git a/src/hamiltonian.jl b/src/hamiltonian.jl index 3659677..e593fa5 100644 --- a/src/hamiltonian.jl +++ b/src/hamiltonian.jl @@ -185,8 +185,10 @@ function solve_ECG( threshold::Real = 0.95, max_attempts::Int = 10 * n, max_condition::Real = 1.0e12, - verbose::Bool = true + verbose::Bool = true, + state::Int = 1 ) + state >= 1 || throw(ArgumentError("state must be >= 1, got $state")) b₁ = float(scale) basis_fns = Rank0Gaussian[] @@ -205,6 +207,7 @@ function solve_ECG( n_accepted = 0 n_rejected = 0 attempt = 0 + E_target_last = Inf # last accepted energy of the target state specifically while n_accepted < n && attempt < max_attempts attempt += 1 @@ -281,14 +284,18 @@ function solve_ECG( continue end - E0 = minimum(λs) - - # Variational principle: adding any linearly independent function to the - # basis cannot raise the ground-state energy. If it does, the candidate - # is numerically near-degenerate with the existing basis (not caught by - # the overlap / condition-number checks above), so reject it. - if n_accepted > 0 && E0 > E_hist[end] + 1.0e-10 - @warn "Candidate raises energy at step $ki, rejecting" ΔE = E0 - E_hist[end] + # Target eigenvalue: use the requested state when available, otherwise + # fall back to the highest available eigenvalue during early build-up. + target_idx = min(state, length(λs)) + E0 = λs[target_idx] + + # Variational principle: the k-th eigenvalue is an upper bound to the + # k-th exact energy, so adding a linearly independent function cannot + # raise it. Only compare against a previous value of the SAME eigenvalue + # (target_idx == state) to avoid spurious rejections during build-up + # when transitioning from tracking a lower eigenvalue to the target one. + if target_idx == state && isfinite(E_target_last) && E0 > E_target_last + 1.0e-10 + @warn "Candidate raises energy at step $ki, rejecting" ΔE = E0 - E_target_last n_rejected += 1 continue end @@ -297,6 +304,9 @@ function solve_ECG( n_accepted += 1 push!(E_hist, E0) push!(vecs_list, Us) + if target_idx == state + E_target_last = E0 + end verbose && @info "Step $n_accepted" E₀ = E0 attempts = attempt rejected = n_rejected end @@ -305,6 +315,186 @@ function solve_ECG( end Emin = last(E_hist) - @info "Optimization complete" E₀ = Emin n_basis = n_accepted - return SolverResults(basis_fns, n_accepted, operators, method, sampler, b₁, Emin, E_hist, vecs_list, E_hist) + @info "Optimization complete" E₀ = Emin n_basis = n_accepted state = state + return SolverResults(basis_fns, n_accepted, operators, method, sampler, b₁, Emin, state, E_hist, vecs_list, E_hist) +end + +# --------------------------------------------------------------------------- +# Operators — ITensors-style accumulator for building a Hamiltonian +# --------------------------------------------------------------------------- + +""" + Operators + +Accumulates `KineticOperator` and `CoulombOperator` terms to build a Hamiltonian. + +# Constructors + + Operators() # system-unaware; add pre-built operators with `+=` + Operators(masses) # system-aware; enables string/index shorthand + +# System-aware interface (requires `Operators(masses)`) + +```julia +ops = Operators([m₁, m₂, m₃]) +ops += "Kinetic" +ops += "Coulomb", 1, 2, +1.0 # pair (1,2) with coupling coefficient +1.0 +ops += "Coulomb", 1, 3, -1.0 +``` + +Particle indices follow the original ordering of `masses`. +The Jacobi transform is handled internally. + +# System-unaware interface (pre-built operators) + +```julia +ops = Operators() +ops += KineticOperator(Λmat) +ops += CoulombOperator(-1.0, w) +``` + +Both interfaces can be mixed freely. Pass `ops` directly to `solve_ECG`, +`solve_ECG_variational`, `solve_ECG_sequential`, or `build_hamiltonian_matrix`. +""" +mutable struct Operators + terms::Vector{FewBodyHamiltonians.Operator} + masses::Union{Nothing, Vector{Float64}} + charges::Union{Nothing, Vector{Float64}} + _U::Union{Nothing, Matrix{Float64}} +end + +Operators() = Operators(FewBodyHamiltonians.Operator[], nothing, nothing, nothing) + +function Operators(masses::Vector{<:Real}) + _, U = _jacobi_transform(Float64.(masses)) + return Operators(FewBodyHamiltonians.Operator[], Float64.(masses), nothing, U) +end + +""" + Operators(masses, charges) + +Create an `Operators` for a system with given particle `masses` and `charges`. +Enables the fully automatic shorthand `ops += "Coulomb"`, which adds all +``N(N-1)/2`` pairwise Coulomb interactions with coefficients ``q_i q_j``. + +```julia +# H⁻: proton (charge +1) + two electrons (charge -1) +ops = Operators([1e15, 1.0, 1.0], [+1, -1, -1]) +ops += "Kinetic" +ops += "Coulomb" # adds (1,2)→-1, (1,3)→-1, (2,3)→+1 automatically +``` +""" +function Operators(masses::Vector{<:Real}, charges::Vector{<:Real}) + length(masses) == length(charges) || + throw(ArgumentError("masses and charges must have the same length")) + _, U = _jacobi_transform(Float64.(masses)) + return Operators(FewBodyHamiltonians.Operator[], Float64.(masses), Float64.(charges), U) +end + +function Base.:+(ops::Operators, op::FewBodyHamiltonians.Operator) + push!(ops.terms, op) + return ops +end + +function Base.:+(ops::Operators, name::AbstractString) + if name == "Kinetic" + ops.masses !== nothing || + throw(ArgumentError("\"Kinetic\" requires Operators(masses).")) + push!(ops.terms, KineticOperator(ops.masses)) + elseif name == "Coulomb" + ops.charges !== nothing || + throw(ArgumentError( + "\"Coulomb\" without indices requires Operators(masses, charges). " * + "Use ops += \"Coulomb\", i, j, coeff for explicit pairs." + )) + N = length(ops.masses) + for i in 1:N, j in (i + 1):N + e_ij = zeros(Float64, N) + e_ij[i] = 1.0 + e_ij[j] = -1.0 + w = ops._U' * e_ij + push!(ops.terms, CoulombOperator(ops.charges[i] * ops.charges[j], w)) + end + else + throw(ArgumentError( + "Unknown operator \"$name\". Supported: \"Kinetic\", \"Coulomb\"." + )) + end + return ops +end + +function Base.:+(ops::Operators, term::Tuple{<:AbstractString, <:Integer, <:Integer, <:Real}) + name, i, j, coeff = term + name == "Coulomb" || + throw(ArgumentError("Unknown operator \"$name\". Supported: \"Coulomb\".")) + ops.masses !== nothing || + throw(ArgumentError( + "String-based \"Coulomb\" requires Operators(masses)." + )) + i != j || throw(ArgumentError("Particle indices must be distinct, got i = j = $i.")) + N = length(ops.masses) + 1 ≤ i ≤ N || throw(ArgumentError("Particle index i=$i out of range [1, $N].")) + 1 ≤ j ≤ N || throw(ArgumentError("Particle index j=$j out of range [1, $N].")) + e_ij = zeros(Float64, N) + e_ij[i] = 1.0 + e_ij[j] = -1.0 + w = ops._U' * e_ij + push!(ops.terms, CoulombOperator(Float64(coeff), w)) + return ops +end + +Base.length(ops::Operators) = length(ops.terms) +Base.iterate(ops::Operators) = iterate(ops.terms) +Base.iterate(ops::Operators, state) = iterate(ops.terms, state) +Base.getindex(ops::Operators, i::Int) = ops.terms[i] +Base.eltype(::Type{Operators}) = FewBodyHamiltonians.Operator + +function Base.show(io::IO, ops::Operators) + n = length(ops.terms) + if ops.masses !== nothing && ops.charges !== nothing + header = "Operators(masses=$(round.(ops.masses; sigdigits=3)), charges=$(ops.charges))" + elseif ops.masses !== nothing + header = "Operators($(length(ops.masses))-body)" + else + header = "Operators" + end + println(io, "$header with $n term$(n == 1 ? "" : "s"):") + for op in ops.terms + if op isa KineticOperator + println(io, " + Kinetic") + elseif op isa CoulombOperator + println(io, " + $(op.coefficient) × Coulomb(w = $(round.(op.w; digits = 3)))") + else + println(io, " + $(typeof(op))") + end + end +end + +""" + coulomb_weights(ops::Operators) -> Vector{Vector{Float64}} + +Return the Jacobi-frame weight vectors for every `CoulombOperator` in `ops`, +in the order they were added. Useful for manual basis construction: + +```julia +w_jac = coulomb_weights(ops) +A = _generate_A_matrix(bij, w_jac) +``` +""" +coulomb_weights(ops::Operators) = [op.w for op in ops.terms if op isa CoulombOperator] + +function build_hamiltonian_matrix(basis::BasisSet{<:GaussianBase}, ops::Operators) + return build_hamiltonian_matrix(basis, ops.terms) +end + +function solve_ECG(ops::Operators, n::Int = 50; kwargs...) + return solve_ECG(ops.terms, n; kwargs...) +end + +function solve_ECG_variational(ops::Operators, n::Int = 50; kwargs...) + return solve_ECG_variational(ops.terms, n; kwargs...) +end + +function solve_ECG_sequential(ops::Operators, n::Int = 50; kwargs...) + return solve_ECG_sequential(ops.terms, n; kwargs...) end diff --git a/src/utils.jl b/src/utils.jl index fe03c40..c8db769 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -11,9 +11,10 @@ Container returned by all ECG solvers ([`solve_ECG`](@ref), - `method` : solver symbol (`:quasirandom`, `:random`, `:variational`, `:sequential`). - `sampler` : quasi-/pseudo-random sampler used for basis generation. - `length_scale` : Gaussian width scale passed at construction. -- `ground_state` : lowest eigenvalue (ground-state energy in a.u.). +- `ground_state` : energy of the target eigenstate (ground state by default). +- `state` : which eigenstate was targeted (1 = ground state, 2 = first excited, …). - `energies` : energy after each accepted basis function (stochastic) or after each step (sequential). -- `eigenvectors` : list of eigenvector matrices; `eigenvectors[end][:, 1]` is the ground-state coefficient vector. +- `eigenvectors` : list of eigenvector matrices; `eigenvectors[end][:, state]` is the target-state coefficient vector. - `fg_history` : monotone-decreasing objective value after each gradient evaluation (variational solvers). """ struct SolverResults @@ -24,6 +25,7 @@ struct SolverResults sampler::QuasiMonteCarlo.DeterministicSamplingAlgorithm length_scale::Float64 ground_state::Float64 + state::Int energies::Vector{Float64} eigenvectors::Vector{Matrix{Float64}} fg_history::Vector{Float64} @@ -48,7 +50,7 @@ function ψ₀(r::AbstractVector, c::AbstractVector, basis_fns::Vector{<:Gaussia ) end -function ψ₀(r::AbstractVector, sr::SolverResults; state::Int = 1) +function ψ₀(r::AbstractVector, sr::SolverResults; state::Int = sr.state) c = sr.eigenvectors[end][:, state] return ψ₀(r, c, sr.basis_functions) end diff --git a/src/variational.jl b/src/variational.jl index 3366898..0b90362 100644 --- a/src/variational.jl +++ b/src/variational.jl @@ -329,6 +329,7 @@ function solve_ECG_variational( HaltonSample(), # placeholder: no stochastic sampler is used float(scale), ground_state, + 1, [ground_state], # single-point; no greedy build-up history [evecs], fg_history, @@ -564,6 +565,7 @@ function solve_ECG_sequential( HaltonSample(), float(scale), ground_state, + 1, E_history, # one energy per sequential step — use convergence() [evecs], fg_history, diff --git a/test/test_utils.jl b/test/test_utils.jl index bd5a396..165135f 100644 --- a/test/test_utils.jl +++ b/test/test_utils.jl @@ -44,6 +44,7 @@ function create_mock_solver_results(; HaltonSample(), scale, energies[end], + 1, energies, eigenvectors, energies # fg_history mirrors energies for mock/stochastic results @@ -79,8 +80,8 @@ end basis_fns = [Rank0Gaussian([1.0;;], [0.0])] ops = Operator[KineticOperator([0.5;;])] - sr_halton = SolverResults(basis_fns, 1, ops, :quasirandom, HaltonSample(), 1.0, -0.5, [-0.5], [ones(1, 1)], [-0.5]) - sr_sobol = SolverResults(basis_fns, 1, ops, :quasirandom, SobolSample(), 1.0, -0.5, [-0.5], [ones(1, 1)], [-0.5]) + sr_halton = SolverResults(basis_fns, 1, ops, :quasirandom, HaltonSample(), 1.0, -0.5, 1, [-0.5], [ones(1, 1)], [-0.5]) + sr_sobol = SolverResults(basis_fns, 1, ops, :quasirandom, SobolSample(), 1.0, -0.5, 1, [-0.5], [ones(1, 1)], [-0.5]) @test sr_halton.sampler isa HaltonSample @test sr_sobol.sampler isa SobolSample @@ -390,7 +391,7 @@ end sr = SolverResults( basis_fns, 1, ops, :quasirandom, HaltonSample(), - 1.0, -0.5, [-0.5], eigvecs, [-0.5] + 1.0, -0.5, 1, [-0.5], eigvecs, [-0.5] ) # All utilities should work From 0ae874e75c6c427fdda1a1641b1f40545b8b6812 Mon Sep 17 00:00:00 2001 From: MartinMikkelsen Date: Fri, 6 Mar 2026 13:52:15 +0100 Subject: [PATCH 8/8] updated tests and documentation --- Examples/Variational.jl | 26 +-- docs/src/API.md | 2 + docs/src/examples.md | 78 +++----- src/hamiltonian.jl | 37 ++-- test/runtests.jl | 1 + test/test_operators.jl | 410 ++++++++++++++++++++++++++++++++++++++++ 6 files changed, 466 insertions(+), 88 deletions(-) create mode 100644 test/test_operators.jl diff --git a/Examples/Variational.jl b/Examples/Variational.jl index 1ac3306..40dcdfb 100644 --- a/Examples/Variational.jl +++ b/Examples/Variational.jl @@ -21,21 +21,13 @@ using FewBodyECG using LinearAlgebra -using QuasiMonteCarlo import FewBodyECG: default_scale, BasisSet, Rank0Gaussian masses_Hm = [1.0e15, 1.0, 1.0] # fixed nucleus + 2 electrons -Λ_Hm = Λ(masses_Hm) -_, U_Hm = _jacobi_transform(masses_Hm) -w_pairs = [[1, -1, 0], [1, 0, -1], [0, 1, -1]] -w_raw_Hm = [U_Hm' * Float64.(w) for w in w_pairs] -coeffs_Hm = [-1.0, -1.0, +1.0] # e-nucleus (×2) and e-e repulsion - -ops_Hm = Operator[ - KineticOperator(Λ_Hm); - [CoulombOperator(c, w) for (c, w) in zip(coeffs_Hm, w_raw_Hm)]... -] +ops_Hm = Operators(masses_Hm, [+1, -1, -1]) # proton, e₁, e₂ +ops_Hm += "Kinetic" +ops_Hm += "Coulomb" E_exact_Hm = -0.527751016523 n = 30 @@ -67,16 +59,10 @@ r_grid, ρ = correlation_function(sr_var; rmin = 0.01, rmax = 15.0, npoints = 20 println("\n Correlation function computed: $(length(r_grid)) points, max ρ at r = $(round(r_grid[argmax(ρ)], digits=3)) a.u.") masses_tdμ = [5496.918, 3670.481, 206.7686] # t, d, μ in electron masses -Λ_tdμ = Λ(masses_tdμ) -_, U_tdμ = _jacobi_transform(masses_tdμ) - -w_raw_tdμ = [U_tdμ' * Float64.(w) for w in w_pairs] -coeffs_tdμ = [+1.0, -1.0, -1.0] # t-d repulsion, t-μ and d-μ attraction -ops_tdμ = Operator[ - KineticOperator(Λ_tdμ); - [CoulombOperator(c, w) for (c, w) in zip(coeffs_tdμ, w_raw_tdμ)]... -] +ops_tdμ = Operators(masses_tdμ, [+1, +1, -1]) # triton, deuteron, muon +ops_tdμ += "Kinetic" +ops_tdμ += "Coulomb" # t-d repulsion (+1), t-μ and d-μ attraction (-1) E_exact_tdμ = -111.36444 scale_tdμ = 0.03 # nuclear scale diff --git a/docs/src/API.md b/docs/src/API.md index 84a2407..9da0380 100644 --- a/docs/src/API.md +++ b/docs/src/API.md @@ -11,6 +11,8 @@ solve_ECG_sequential ## Operators ```@docs +Operators +coulomb_weights KineticOperator CoulombOperator ``` diff --git a/docs/src/examples.md b/docs/src/examples.md index 854ee93..2eb35b9 100644 --- a/docs/src/examples.md +++ b/docs/src/examples.md @@ -4,25 +4,13 @@ Suppose you want to calculate the ground state energy of the hydrogen anion in t ```@example example1 using FewBodyECG -using LinearAlgebra using Plots -using QuasiMonteCarlo masses = [1.0e15, 1.0, 1.0] -Λmat = Λ(masses) -kin = KineticOperator(Λmat) -J, U = _jacobi_transform(masses) - -w_list = [[1, -1, 0], [1, 0, -1], [0, 1, -1]] - -w_raw = [U' * w for w in w_list] -coeffs = [-1.0, -1.0, +1.0] - -ops = Operator[ - kin; - (CoulombOperator(c, w) for (c, w) in zip(coeffs, w_raw))... -] +ops = Operators(masses, [+1, -1, -1]) # proton, e₁, e₂ +ops += "Kinetic" +ops += "Coulomb" # adds all 3 pairwise interactions automatically result = solve_ECG(ops, 250, scale = 1.0) @@ -46,20 +34,13 @@ L-BFGS via [OptimKit.jl](https://github.com/Jutho/OptimKit.jl). ```@example example_var_fresh using FewBodyECG -using LinearAlgebra using Plots masses = [1.0e15, 1.0, 1.0] # H⁻: fixed nucleus + 2 electrons -Λmat = Λ(masses) -kin = KineticOperator(Λmat) -J, U = _jacobi_transform(masses) - -w_list = [[1, -1, 0], [1, 0, -1], [0, 1, -1]] -w_raw = [U' * w for w in w_list] -coeffs = [-1.0, -1.0, +1.0] - -ops = Operator[kin; [CoulombOperator(c, w) for (c, w) in zip(coeffs, w_raw)]...] +ops = Operators(masses, [+1, -1, -1]) +ops += "Kinetic" +ops += "Coulomb" sr = solve_ECG_variational(ops, 20; scale = 1.0, max_iterations = 200, verbose = false) @@ -85,15 +66,13 @@ Here we apply it to the hydrogen atom ground state (1s): ```@example example_seq using FewBodyECG -using LinearAlgebra using Plots -masses = [1.0e15, 1.0] # hydrogen: heavy nucleus + electron -Λmat = Λ(masses) -_, U = _jacobi_transform(masses) -w = U' * Float64.([1, -1]) # electron–nucleus separation +masses = [1.0e15, 1.0] # hydrogen: heavy nucleus + electron -ops = Operator[KineticOperator(Λmat); CoulombOperator(-1.0, w)] +ops = Operators(masses) +ops += "Kinetic" +ops += ("Coulomb", 1, 2, -1.0) # electron–nucleus attraction sr = solve_ECG_sequential(ops, 12; n_candidates = 8, scale = 1.0, max_iterations_step = 80, verbose = false) @@ -126,10 +105,9 @@ using LinearAlgebra using Plots masses = [1.0e15, 1.0] -Λmat = Λ(masses) -_, U = _jacobi_transform(masses) -w = U' * Float64.([1, -1]) -ops = Operator[KineticOperator(Λmat); CoulombOperator(-1.0, w)] +ops = Operators(masses) +ops += "Kinetic" +ops += ("Coulomb", 1, 2, -1.0) a_p = [1.0] # polarisation along the single Jacobi coordinate s_zero = [0.0] @@ -169,10 +147,9 @@ using LinearAlgebra using Plots masses = [1.0e15, 1.0] -Λmat = Λ(masses) -_, U = _jacobi_transform(masses) -w = U' * Float64.([1, -1]) -ops = Operator[KineticOperator(Λmat); CoulombOperator(-1.0, w)] +ops = Operators(masses) +ops += "Kinetic" +ops += ("Coulomb", 1, 2, -1.0) # Orthogonal Cartesian directions → pure d-wave channel (a ⊥ b) a_d = reshape([1.0, 0.0, 0.0], 1, 3) @@ -212,21 +189,13 @@ be passed directly via the `scale` keyword. ```@example example_tdmu using FewBodyECG -using LinearAlgebra using Plots masses = [5496.918, 3670.481, 206.7686] # t, d, μ in electron masses -Λmat = Λ(masses) -_, U = _jacobi_transform(masses) -w_pairs = [[1, -1, 0], [1, 0, -1], [0, 1, -1]] -w_raw = [U' * Float64.(w) for w in w_pairs] -coeffs = [+1.0, -1.0, -1.0] # t-d repulsion; t-μ and d-μ attraction - -ops = Operator[ - KineticOperator(Λmat); - [CoulombOperator(c, w) for (c, w) in zip(coeffs, w_raw)]... -] +ops = Operators(masses, [+1, +1, -1]) # triton, deuteron, muon +ops += "Kinetic" +ops += "Coulomb" # t-d repulsion (+1), t-μ and d-μ attraction (-1) result = solve_ECG(ops, 100; scale = 0.03, verbose = false) @@ -248,13 +217,12 @@ fresh L-BFGS run. ```@example example_var_warm using FewBodyECG -using LinearAlgebra masses = [1.0e15, 1.0, 1.0] -Λmat = Λ(masses); kin = KineticOperator(Λmat) -J, U = _jacobi_transform(masses) -w_raw = [U' * w for w in [[1,-1,0],[1,0,-1],[0,1,-1]]] -ops = Operator[kin; [CoulombOperator(c,w) for (c,w) in zip([-1.,-1.,1.], w_raw)]...] + +ops = Operators(masses, [+1, -1, -1]) +ops += "Kinetic" +ops += "Coulomb" sr_stoch = solve_ECG(ops, 20; scale = 1.0, verbose = false) basis0 = BasisSet(Rank0Gaussian[sr_stoch.basis_functions...]) diff --git a/src/hamiltonian.jl b/src/hamiltonian.jl index e593fa5..5b9a8d3 100644 --- a/src/hamiltonian.jl +++ b/src/hamiltonian.jl @@ -319,31 +319,40 @@ function solve_ECG( return SolverResults(basis_fns, n_accepted, operators, method, sampler, b₁, Emin, state, E_hist, vecs_list, E_hist) end -# --------------------------------------------------------------------------- -# Operators — ITensors-style accumulator for building a Hamiltonian -# --------------------------------------------------------------------------- - """ Operators Accumulates `KineticOperator` and `CoulombOperator` terms to build a Hamiltonian. +Handles Jacobi coordinate transforms internally so that callers can work +with physical particle indices rather than Jacobi-frame weight vectors. # Constructors - Operators() # system-unaware; add pre-built operators with `+=` - Operators(masses) # system-aware; enables string/index shorthand + Operators() # system-unaware; add pre-built operators with `+=` + Operators(masses) # system-aware; enables string/index shorthand + Operators(masses, charges) # fully automatic; enables `ops += "Coulomb"` shorthand + +# System-aware interface -# System-aware interface (requires `Operators(masses)`) +Particle indices follow the original ordering of `masses`. All Jacobi +transforms are computed internally. ```julia ops = Operators([m₁, m₂, m₃]) ops += "Kinetic" -ops += "Coulomb", 1, 2, +1.0 # pair (1,2) with coupling coefficient +1.0 -ops += "Coulomb", 1, 3, -1.0 +ops += ("Coulomb", 1, 2, +1.0) # pair (1,2) with coupling coefficient +1.0 +ops += ("Coulomb", 1, 3, -1.0) ``` -Particle indices follow the original ordering of `masses`. -The Jacobi transform is handled internally. +When charges are also supplied, the fully-automatic shorthand `ops += "Coulomb"` +adds all ``N(N-1)/2`` pairwise terms with coefficients ``q_i q_j``: + +```julia +# Helium atom: nucleus (Z=2), two electrons +ops = Operators([1e15, 1.0, 1.0], [+2, -1, -1]) +ops += "Kinetic" +ops += "Coulomb" # adds (1,2)→-2, (1,3)→-2, (2,3)→+1 automatically +``` # System-unaware interface (pre-built operators) @@ -353,8 +362,10 @@ ops += KineticOperator(Λmat) ops += CoulombOperator(-1.0, w) ``` -Both interfaces can be mixed freely. Pass `ops` directly to `solve_ECG`, -`solve_ECG_variational`, `solve_ECG_sequential`, or `build_hamiltonian_matrix`. +Both interfaces can be mixed freely. Pass `ops` directly to [`solve_ECG`](@ref), +[`solve_ECG_variational`](@ref), [`solve_ECG_sequential`](@ref), or +`build_hamiltonian_matrix`. Use [`coulomb_weights`](@ref) to retrieve +the Jacobi-frame weight vectors for manual basis construction. """ mutable struct Operators terms::Vector{FewBodyHamiltonians.Operator} diff --git a/test/runtests.jl b/test/runtests.jl index 77e882b..951d759 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -13,5 +13,6 @@ using FewBodyECG include("test_utils.jl") include("test_types.jl") include("test_variational.jl") + include("test_operators.jl") end diff --git a/test/test_operators.jl b/test/test_operators.jl new file mode 100644 index 0000000..f870dc8 --- /dev/null +++ b/test/test_operators.jl @@ -0,0 +1,410 @@ +using Test +using LinearAlgebra +using FewBodyHamiltonians +using FewBodyECG +import FewBodyECG: _jacobi_transform, Λ + +@testset "Operators" begin + + @testset "Construction" begin + + @testset "Empty (system-unaware)" begin + ops = Operators() + @test length(ops) == 0 + @test ops.masses === nothing + @test ops.charges === nothing + @test ops._U === nothing + end + + @testset "With masses only" begin + masses = [1.0e15, 1.0] + ops = Operators(masses) + @test ops.masses ≈ Float64.(masses) + @test ops.charges === nothing + @test ops._U !== nothing + @test size(ops._U) == (2, 1) # N × (N-1) Jacobi matrix + @test length(ops) == 0 + end + + @testset "With masses and charges" begin + masses = [1.0e15, 1.0, 1.0] + charges = [+1, -1, -1] + ops = Operators(masses, charges) + @test ops.masses ≈ Float64.(masses) + @test ops.charges ≈ Float64.(charges) + @test ops._U !== nothing + @test size(ops._U) == (3, 2) # N × (N-1) Jacobi matrix + end + + @testset "Error: mismatched masses/charges" begin + @test_throws ArgumentError Operators([1.0, 2.0], [1.0]) + @test_throws ArgumentError Operators([1.0, 2.0, 3.0], [1.0, -1.0]) + end + end + + @testset "Adding raw operators" begin + + @testset "KineticOperator" begin + ops = Operators() + K = KineticOperator([0.5;;]) + ops += K + @test length(ops) == 1 + @test ops[1] isa KineticOperator + end + + @testset "CoulombOperator" begin + ops = Operators() + V = CoulombOperator(-1.0, [1.0]) + ops += V + @test length(ops) == 1 + @test ops[1] isa CoulombOperator + @test ops[1].coefficient ≈ -1.0 + end + + @testset "Multiple raw operators" begin + ops = Operators() + ops += KineticOperator([0.5;;]) + ops += CoulombOperator(-1.0, [1.0]) + ops += CoulombOperator(+0.5, [1.0]) + @test length(ops) == 3 + end + end + + @testset "String: Kinetic" begin + + @testset "Adds correct KineticOperator" begin + masses = [1.0e15, 1.0] + ops = Operators(masses) + ops += "Kinetic" + @test length(ops) == 1 + @test ops[1] isa KineticOperator + end + + @testset "Three-body kinetic" begin + masses = [1.0e15, 1.0, 1.0] + ops = Operators(masses) + ops += "Kinetic" + @test length(ops) == 1 + @test ops[1] isa KineticOperator + end + + @testset "Error: no masses" begin + ops = Operators() + @test_throws ArgumentError ops += "Kinetic" + end + + @testset "Error: unknown string" begin + ops = Operators([1.0e15, 1.0]) + @test_throws ArgumentError ops += "BadOp" + end + end + + @testset "Tuple: explicit Coulomb pair" begin + + @testset "Adds correct CoulombOperator" begin + masses = [1.0e15, 1.0, 1.0] + ops = Operators(masses) + ops += ("Coulomb", 1, 2, -1.0) + @test length(ops) == 1 + @test ops[1] isa CoulombOperator + @test ops[1].coefficient ≈ -1.0 + end + + @testset "Coefficient is stored correctly" begin + masses = [1.0e15, 1.0, 1.0] + ops = Operators(masses) + ops += ("Coulomb", 1, 3, +2.5) + @test ops[1].coefficient ≈ 2.5 + end + + @testset "Weight vector has correct dimension" begin + masses = [1.0e15, 1.0, 1.0] # 3-body → 2 Jacobi dims + ops = Operators(masses) + ops += ("Coulomb", 1, 2, -1.0) + @test length(ops[1].w) == 2 + end + + @testset "Error: no masses" begin + ops = Operators() + @test_throws ArgumentError ops += ("Coulomb", 1, 2, -1.0) + end + + @testset "Error: same index (i == j)" begin + ops = Operators([1.0e15, 1.0, 1.0]) + @test_throws ArgumentError ops += ("Coulomb", 1, 1, -1.0) + @test_throws ArgumentError ops += ("Coulomb", 2, 2, -1.0) + end + + @testset "Error: index out of range" begin + ops = Operators([1.0e15, 1.0, 1.0]) + @test_throws ArgumentError ops += ("Coulomb", 0, 1, -1.0) + @test_throws ArgumentError ops += ("Coulomb", 1, 4, -1.0) + @test_throws ArgumentError ops += ("Coulomb", -1, 2, -1.0) + end + + @testset "Error: unknown operator name in tuple" begin + ops = Operators([1.0e15, 1.0, 1.0]) + @test_throws ArgumentError ops += ("Kinetic", 1, 2, -1.0) + @test_throws ArgumentError ops += ("Unknown", 1, 2, -1.0) + end + end + + @testset "String: auto all-pairs Coulomb" begin + + @testset "Two-body: one pair" begin + ops = Operators([1.0e15, 1.0], [+1, -1]) + ops += "Coulomb" + @test length(ops) == 1 + @test ops[1] isa CoulombOperator + @test ops[1].coefficient ≈ -1.0 # (+1) * (-1) + end + + @testset "Three-body: three pairs" begin + ops = Operators([1.0e15, 1.0, 1.0], [+1, -1, -1]) + ops += "Coulomb" + @test length(ops) == 3 + @test all(op isa CoulombOperator for op in ops) + end + + @testset "Coefficients follow q_i * q_j" begin + # proton (+1), e₁ (-1), e₂ (-1) + ops = Operators([1.0e15, 1.0, 1.0], [+1, -1, -1]) + ops += "Coulomb" + coeffs = [ops[i].coefficient for i in 1:3] + # pairs (1,2), (1,3), (2,3) → -1, -1, +1 + @test coeffs[1] ≈ -1.0 + @test coeffs[2] ≈ -1.0 + @test coeffs[3] ≈ +1.0 + end + + @testset "Four-body: six pairs" begin + ops = Operators([1.0e15, 1.0, 1.0, 1.0], [+1, -1, -1, -1]) + ops += "Coulomb" + @test length(ops) == 6 # C(4,2) + end + + @testset "Error: no charges" begin + ops = Operators([1.0e15, 1.0, 1.0]) + @test_throws ArgumentError ops += "Coulomb" + end + end + + @testset "Collection interface" begin + + @testset "length" begin + ops = Operators([1.0e15, 1.0, 1.0], [+1, -1, -1]) + @test length(ops) == 0 + ops += "Kinetic" + @test length(ops) == 1 + ops += "Coulomb" + @test length(ops) == 4 # 1 kinetic + 3 Coulomb + end + + @testset "getindex" begin + ops = Operators([1.0e15, 1.0, 1.0], [+1, -1, -1]) + ops += "Kinetic" + ops += "Coulomb" + @test ops[1] isa KineticOperator + @test ops[2] isa CoulombOperator + @test ops[4] isa CoulombOperator + end + + @testset "iterate" begin + ops = Operators([1.0e15, 1.0, 1.0], [+1, -1, -1]) + ops += "Kinetic" + ops += "Coulomb" + collected = collect(ops) + @test collected[1] isa KineticOperator + @test all(op isa CoulombOperator for op in collected[2:4]) + end + + @testset "eltype" begin + @test eltype(Operators) == FewBodyHamiltonians.Operator + end + end + + @testset "show" begin + + @testset "Empty (no masses)" begin + ops = Operators() + str = sprint(show, ops) + @test occursin("Operators", str) + @test occursin("0 terms", str) + end + + @testset "With masses (n-body header)" begin + ops = Operators([1.0e15, 1.0]) + ops += "Kinetic" + str = sprint(show, ops) + @test occursin("2-body", str) + @test occursin("Kinetic", str) + end + + @testset "With masses and charges" begin + ops = Operators([1.0e15, 1.0, 1.0], [+1, -1, -1]) + ops += "Kinetic" + ops += "Coulomb" + str = sprint(show, ops) + @test occursin("charges", str) + @test occursin("Kinetic", str) + @test occursin("Coulomb", str) + end + + @testset "Singular 'term' for 1 operator" begin + ops = Operators() + ops += KineticOperator([0.5;;]) + str = sprint(show, ops) + @test occursin("1 term", str) + @test !occursin("1 terms", str) + end + end + + @testset "coulomb_weights" begin + + @testset "Returns only Coulomb weights (skips Kinetic)" begin + ops = Operators([1.0e15, 1.0, 1.0], [+1, -1, -1]) + ops += "Kinetic" + ops += "Coulomb" + ws = coulomb_weights(ops) + @test length(ws) == 3 # 3 Coulomb pairs, not 4 + end + + @testset "Empty if no Coulomb operators" begin + ops = Operators([1.0e15, 1.0]) + ops += "Kinetic" + @test isempty(coulomb_weights(ops)) + end + + @testset "Correct Jacobi dimension" begin + # 2-body → 1 Jacobi coordinate + ops2 = Operators([1.0e15, 1.0], [+1, -1]) + ops2 += "Coulomb" + @test all(length(w) == 1 for w in coulomb_weights(ops2)) + + # 3-body → 2 Jacobi coordinates + ops3 = Operators([1.0e15, 1.0, 1.0], [+1, -1, -1]) + ops3 += "Coulomb" + @test all(length(w) == 2 for w in coulomb_weights(ops3)) + end + + @testset "Returns Vector{Vector{Float64}}" begin + ops = Operators([1.0e15, 1.0], [+1, -1]) + ops += "Coulomb" + ws = coulomb_weights(ops) + @test ws isa Vector + @test eltype(ws) == Vector{Float64} + end + end + + @testset "Equivalence with manual construction" begin + + @testset "Two-body hydrogen atom" begin + masses = [1.0e15, 1.0] + + # New Operators interface + ops_new = Operators(masses) + ops_new += "Kinetic" + ops_new += ("Coulomb", 1, 2, -1.0) + + # Manual interface + Λmat = Λ(masses) + _, U = _jacobi_transform(masses) + w = U' * [1.0, -1.0] + ops_old = Operator[KineticOperator(Λmat); CoulombOperator(-1.0, w)] + + g = Rank0Gaussian([1.0;;], [0.0]) + basis = BasisSet([g]) + + H_new = build_hamiltonian_matrix(basis, ops_new) + H_old = build_hamiltonian_matrix(basis, ops_old) + @test H_new ≈ H_old rtol = 1.0e-12 + end + + @testset "Three-body H⁻ with auto-Coulomb" begin + masses = [1.0e15, 1.0, 1.0] + charges = [+1, -1, -1] + + ops_new = Operators(masses, charges) + ops_new += "Kinetic" + ops_new += "Coulomb" + + Λmat = Λ(masses) + _, U = _jacobi_transform(masses) + w_list = [U' * Float64.(w) for w in [[1, -1, 0], [1, 0, -1], [0, 1, -1]]] + coeffs = [-1.0, -1.0, +1.0] + ops_old = Operator[ + KineticOperator(Λmat); + [CoulombOperator(c, w) for (c, w) in zip(coeffs, w_list)]... + ] + + g = Rank0Gaussian([1.0 0.0; 0.0 1.0], [0.0, 0.0]) + basis = BasisSet([g]) + + H_new = build_hamiltonian_matrix(basis, ops_new) + H_old = build_hamiltonian_matrix(basis, ops_old) + @test H_new ≈ H_old rtol = 1.0e-12 + end + + @testset "Explicit pair Coulomb matches manual" begin + masses = [1.0e15, 1.0, 1.0] + _, U = _jacobi_transform(masses) + + ops_new = Operators(masses) + ops_new += ("Coulomb", 1, 2, -1.0) + ops_new += ("Coulomb", 1, 3, -1.0) + ops_new += ("Coulomb", 2, 3, +1.0) + + w12 = U' * [1.0, -1.0, 0.0] + w13 = U' * [1.0, 0.0, -1.0] + w23 = U' * [0.0, 1.0, -1.0] + ops_old = Operator[ + CoulombOperator(-1.0, w12), + CoulombOperator(-1.0, w13), + CoulombOperator(+1.0, w23), + ] + + g = Rank0Gaussian([1.0 0.0; 0.0 1.0], [0.0, 0.0]) + basis = BasisSet([g]) + + H_new = build_hamiltonian_matrix(basis, ops_new) + H_old = build_hamiltonian_matrix(basis, ops_old) + @test H_new ≈ H_old rtol = 1.0e-12 + end + end + + @testset "Integration with solvers" begin + + @testset "solve_ECG: hydrogen atom ≈ -0.5 Ha" begin + masses = [1.0e15, 1.0] + ops = Operators(masses) + ops += "Kinetic" + ops += ("Coulomb", 1, 2, -1.0) + sr = solve_ECG(ops, 25; scale = 1.0, verbose = false) + @test sr.ground_state < -0.46 # converging toward -0.5 Ha + @test sr.ground_state > -0.52 + end + + @testset "solve_ECG: H⁻ auto-Coulomb (bound state)" begin + masses = [1.0e15, 1.0, 1.0] + ops = Operators(masses, [+1, -1, -1]) + ops += "Kinetic" + ops += "Coulomb" + sr = solve_ECG(ops, 15; scale = 1.0, verbose = false) + # H⁻ ground state ≈ -0.528 Ha; with a small basis just verify it is bound + @test sr.ground_state < -0.3 + end + + @testset "solve_ECG via Operators matches ops.terms" begin + masses = [1.0e15, 1.0] + ops = Operators(masses) + ops += "Kinetic" + ops += ("Coulomb", 1, 2, -1.0) + + sr_ops = solve_ECG(ops, 25; scale = 1.0, verbose = false) + sr_vec = solve_ECG(ops.terms, 25; scale = 1.0, verbose = false) + + @test sr_ops.ground_state < -0.46 + @test sr_vec.ground_state < -0.46 + end + end +end