BFwf with many grads and laps + Hyperduals

.github/workflows/CI.yml

@@ -18,8 +18,8 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.0'
           - '1.8'
+          - '1.9'
           - 'nightly'
         os:
           - ubuntu-latest

Project.toml

@@ -5,25 +5,35 @@ version = "0.0.1"
 
 [deps]
 ACEbase = "14bae519-eb20-449c-a949-9c58ed33163e"
-ACEcore = "44c1e890-45d1-48ea-94d6-c2ea5b573f71"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+HyperDualNumbers = "50ceba7f-c3ee-5a84-a6e8-3ad40456ec97"
 JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
+LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
 ObjectPools = "658cac36-ff0f-48ad-967c-110375d98c9d"
+Octavian = "6fd5a793-0b7e-452c-907f-f8bfe9c57db4"
+Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 Polynomials4ML = "03c4bcba-a943-47e9-bfa1-b1661fc2974f"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
+StrideArrays = "d1fa6d79-ef01-42a6-86c9-f7c551f8593b"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
-ACEcore = "0.0.1"
+ACEbase = "0.4.2"
 BenchmarkTools = "1"
 ForwardDiff = "0.10"
-ObjectPools = "0.0.2"
-Polynomials4ML = "0.0.2"
-julia = "1"
 JSON = "0.21"
-ACEbase = "0.2"
-
+ObjectPools = "0.3.1"
+julia = "1"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

benchmarks/benchmark_bflow_lux.jl

@@ -0,0 +1,58 @@
+using ACEpsi, Polynomials4ML, StaticArrays, Test 
+using Polynomials4ML: natural_indices, degree, SparseProduct
+using ACEpsi.AtomicOrbitals: Nuc, make_nlms_spec, evaluate
+using ACEpsi: BackflowPooling, BFwf_lux, setupBFState, Jastrow
+using ACEpsi.vmc: gradient, laplacian, grad_params
+using ACEbase.Testing: print_tf, fdtest
+using LuxCore
+using Lux
+using Zygote
+using Optimisers # mainly for the destrcuture(ps) function
+using Random
+using Printf
+using LinearAlgebra
+using BenchmarkTools
+using HyperDualNumbers: Hyper
+
+Rnldegree = n1 = 2
+Ylmdegree = 3
+totdegree = 3
+Nel = 2
+X = randn(SVector{3, Float64}, Nel)
+Σ = rand(spins(), Nel)
+nuclei = [ Nuc(3 * rand(SVector{3, Float64}), 1.0) for _=1:3 ]
+
+# wrap it as HyperDualNumbers
+x2dualwrtj(x, j) = SVector{3}([Hyper(x[i], i == j, i == j, 0) for i = 1:3])
+hX = [x2dualwrtj(x, 0) for x in X]
+hX[1] = x2dualwrtj(X[1], 1) # test eval for grad wrt x coord of first elec
+
+##
+
+# Defining AtomicOrbitalsBasis
+n2 = 2
+Pn = Polynomials4ML.legendre_basis(n1+1)
+spec = [(n1 = n1, n2 = n2, l = l) for n1 = 1:n1 for n2 = 1:n2 for l = 0:n1-1] 
+ζ = rand(length(spec))
+Dn = GaussianBasis(ζ)
+bRnl = AtomicOrbitalsRadials(Pn, Dn, spec) 
+bYlm = RYlmBasis(Ylmdegree)
+
+# setup state
+BFwf_chain, spec, spec1p = BFwf_lux(Nel, bRnl, bYlm, nuclei; totdeg = totdegree, ν = 2)
+ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ)
+
+##
+
+@info("Test evaluate")
+
+@btime BFwf_chain($X, $ps, $st)
+@btime gradient($BFwf_chain, $X, $ps, $st)
+@btime laplacian($BFwf_chain, $X, $ps, $st)
+
+
+@profview let  BFwf_chain = BFwf_chain, X = X, ps =  ps, st = st
+   for i = 1:10_000
+      BFwf_chain(X, ps, st)
+   end
+end

examples/3D/Be.jl

@@ -0,0 +1,65 @@
+using ACEpsi, Polynomials4ML, StaticArrays, Test 
+using Polynomials4ML: natural_indices, degree, SparseProduct
+using ACEpsi.AtomicOrbitals: Nuc, make_nlms_spec, evaluate
+using ACEpsi: BackflowPooling, BFwf_lux, setupBFState, Jastrow
+using ACEpsi.vmc: gradient, laplacian, grad_params, SumH, MHSampler, VMC, gd_GradientByVMC
+using ACEbase.Testing: print_tf, fdtest
+using LuxCore
+using Lux
+using Zygote
+using Optimisers # mainly for the destrcuture(ps) function
+using Random
+using Printf
+using LinearAlgebra
+using BenchmarkTools
+using HyperDualNumbers: Hyper
+
+n1 = Rnldegree = 5
+Ylmdegree = 2
+totdegree = 5
+Nel = 4
+X = randn(SVector{3, Float64}, Nel)
+Σ = [↑,↑,↓,↓]
+nuclei = [ Nuc(zeros(SVector{3, Float64}), Nel * 1.0)]
+##
+
+# Defining AtomicOrbitalsBasis
+n2 = 1
+Pn = Polynomials4ML.legendre_basis(n1+1)
+spec = [(n1 = n1, n2 = n2, l = l) for n1 = 1:n1 for n2 = 1:n2 for l = 0:n1-1] 
+ζ = 10 * rand(length(spec))
+Dn = SlaterBasis(ζ)
+bRnl = AtomicOrbitalsRadials(Pn, Dn, spec) 
+bYlm = RYlmBasis(Ylmdegree)
+
+# setup state
+wf, spec, spec1p = BFwf_chain, spec, spec1p  = BFwf_lux(Nel, bRnl, bYlm, nuclei; totdeg = totdegree, ν = 2)
+displayspec(spec, spec1p)
+
+ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ)
+p, = destructure(ps)
+length(p)
+
+K(wf, X::AbstractVector, ps, st) = -0.5 * laplacian(wf, X, ps, st)
+Vext(wf, X::AbstractVector, ps, st) = -sum(nuclei[i].charge/norm(nuclei[i].rr - X[j]) for i = 1:length(nuclei) for j in 1:length(X))
+Vee(wf, X::AbstractVector, ps, st) = sum(1/norm(X[i]-X[j]) for i = 1:length(X)-1 for j = i+1:length(X))
+
+ham = SumH(K, Vext, Vee)
+sam = MHSampler(wf, Nel, Δt = 0.5, burnin = 1000, nchains = 2000)
+
+opt_vmc = VMC(3000, 0.1, ACEpsi.vmc.adamW(), lr_dc = 100)
+wf, err_opt, ps = gd_GradientByVMC(opt_vmc, sam, ham, wf, ps, st)
+
+err = err_opt
+per = 0.2
+err1 = zero(err)
+for i = 1:length(err)
+    err1[i] = mean(err[Int(ceil(i-per  * i)):i])
+end
+err1
+
+Eref = -14.667
+
+
+# using Plots
+# plot(abs.(err1 .- Eref), w = 3, yscale=:log10)

examples/3D/He.jl

@@ -0,0 +1,65 @@
+using ACEpsi, Polynomials4ML, StaticArrays, Test 
+using Polynomials4ML: natural_indices, degree, SparseProduct
+using ACEpsi.AtomicOrbitals: Nuc, make_nlms_spec, evaluate
+using ACEpsi: BackflowPooling, BFwf_lux, setupBFState, Jastrow, displayspec
+using ACEpsi.vmc: gradient, laplacian, grad_params, SumH, MHSampler, VMC, gd_GradientByVMC, AdamW, SR
+using ACEbase.Testing: print_tf, fdtest
+using LuxCore
+using Lux
+using Zygote
+using Optimisers # mainly for the destrcuture(ps) function
+using Random
+using Printf
+using LinearAlgebra
+using BenchmarkTools
+using HyperDualNumbers: Hyper
+
+n1 = Rnldegree = 4
+Ylmdegree = 2
+totdegree = 4
+Nel = 2
+X = randn(SVector{3, Float64}, Nel)
+Σ = [↑,↓]
+nuclei = [ Nuc(zeros(SVector{3, Float64}), 2.0)]
+##
+
+# Defining AtomicOrbitalsBasis
+n2 = 1
+Pn = Polynomials4ML.legendre_basis(n1+1)
+spec = [(n1 = n1, n2 = n2, l = l) for n1 = 1:n1 for n2 = 1:n2 for l = 0:n1-1] 
+ζ = 10 * rand(length(spec))
+Dn = SlaterBasis(ζ)
+bRnl = AtomicOrbitalsRadials(Pn, Dn, spec) 
+bYlm = RYlmBasis(Ylmdegree)
+
+# setup state
+wf, spec, spec1p = BFwf_chain, spec, spec1p  = BFwf_lux(Nel, bRnl, bYlm, nuclei; totdeg = totdegree, ν = 2)
+displayspec(spec, spec1p)
+
+ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ)
+p, = destructure(ps)
+length(p)
+
+K(wf, X::AbstractVector, ps, st) = -0.5 * laplacian(wf, X, ps, st)
+Vext(wf, X::AbstractVector, ps, st) = -sum(nuclei[i].charge/norm(nuclei[i].rr - X[j]) for i = 1:length(nuclei) for j in 1:length(X))
+Vee(wf, X::AbstractVector, ps, st) = sum(1/norm(X[i]-X[j]) for i = 1:length(X)-1 for j = i+1:length(X))
+
+ham = SumH(K, Vext, Vee)
+sam = MHSampler(wf, Nel, Δt = 0.5, burnin = 1000, nchains = 2000)
+
+opt_vmc = VMC(3000, 0.0015, SR(1e-4, 0.015), lr_dc = 50)
+wf, err_opt, ps = gd_GradientByVMC(opt_vmc, sam, ham, wf, ps, st)
+
+
+
+err = err_opt
+per = 0.2
+err1 = zero(err)
+for i = 1:length(err)
+    err1[i] = mean(err[Int(ceil(i-per  * i)):i])
+end
+
+Eref = -2.9037
+
+# using Plots
+# plot(abs.(err1 .- Eref), w = 3, yscale=:log10)

examples/3D/harmonic_oscillator.jl

@@ -0,0 +1,53 @@
+using ACEpsi, Polynomials4ML, StaticArrays, Test 
+using Polynomials4ML: natural_indices, degree, SparseProduct
+using ACEpsi.AtomicOrbitals: Nuc, make_nlms_spec, evaluate
+using ACEpsi: BackflowPooling, BFwf_lux, setupBFState, Jastrow, displayspec
+using ACEpsi.vmc: gradient, laplacian, grad_params, SumH, MHSampler, VMC, gd_GradientByVMC
+using ACEbase.Testing: print_tf, fdtest
+using LuxCore
+using Lux
+using Zygote
+using Optimisers # mainly for the destrcuture(ps) function
+using Random
+using Printf
+using LinearAlgebra
+using BenchmarkTools
+using HyperDualNumbers: Hyper
+
+
+
+n1 = Rnldegree = 6
+Ylmdegree = 1
+totdegree = 10
+Nel = 3
+X = randn(SVector{3, Float64}, Nel)
+Σ = [↑,↑,↓]
+nuclei = [ Nuc(zeros(SVector{3, Float64}), 1.0)]
+##
+
+# Defining AtomicOrbitalsBasis
+n2 = 1
+Pn = Polynomials4ML.legendre_basis(n1+1)
+spec = [(n1 = n1, n2 = n2, l = l) for n1 = 1:n1 for n2 = 1:n2 for l = 0:n1-1] 
+ζ = rand(length(spec))
+Dn = GaussianBasis(ζ)
+bRnl = AtomicOrbitalsRadials(Pn, Dn, spec) 
+bYlm = RYlmBasis(Ylmdegree)
+
+# setup state
+wf, spec, spec1p = BFwf_chain, spec, spec1p  = BFwf_lux(Nel, bRnl, bYlm, nuclei; totdeg = totdegree, ν = 2)
+ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ)
+
+wf(X, ps, st)
+
+K(wf, X::AbstractVector, ps, st) = -0.5 * laplacian(wf, X, ps, st)
+
+Vext(wf, X::AbstractVector, ps, st) = 1/2 * sum([X[i][j]^2 for i = 1:length(X) for j = 1:3]) #sum(1/norm(nuclei[i].rr - X[j]) for i = 1:length(nuclei) for j in 1:length(X))
+
+Vee(wf, X::AbstractVector, ps, st) = 0 # sum(1/norm(X[i]-X[j]) for i = 1:length(X)-1 for j = i+1:length(X))
+
+ham = SumH(K, Vext, Vee)
+sam = MHSampler(wf, Nel, nuclei, Δt = 1.0, burnin = 1000, nchains = 2000)
+
+opt_vmc = VMC(100, 0.05, ACEpsi.vmc.adamW(); lr_dc = 30.0)
+wf, err_opt, ps = gd_GradientByVMC(opt_vmc, sam, ham, wf, ps, st)

examples/3D/sto-3g_H.jl

@@ -0,0 +1,63 @@
+using ACEpsi, Polynomials4ML, StaticArrays, Test 
+using Polynomials4ML: natural_indices, degree, SparseProduct
+using ACEpsi.AtomicOrbitals: Nuc, make_nlms_spec, evaluate
+using ACEpsi: BackflowPooling, BFwf_lux, setupBFState, Jastrow, displayspec
+using ACEpsi.vmc: gradient, laplacian, grad_params, SumH, MHSampler, VMC, gd_GradientByVMC
+using ACEbase.Testing: print_tf, fdtest
+using LuxCore
+using Lux
+using Zygote
+using Optimisers # mainly for the destrcuture(ps) function
+using Random
+using Printf
+using LinearAlgebra
+using BenchmarkTools
+using HyperDualNumbers: Hyper
+
+n1 = Rnldegree = 1
+Ylmdegree = 0
+totdegree = 4
+Nel = 1
+X = randn(SVector{3, Float64}, Nel)
+Σ = [↑]
+nuclei = [Nuc(SVector(0.0,0.0,0.0), 1.0)]
+##
+
+# Defining AtomicOrbitalsBasis
+n2 = 1
+Pn = Polynomials4ML.legendre_basis(n1+1)
+spec = [(n1 = 1, n2 = 1, l = 0) ] 
+
+ζ = [0.3425250914E+01, 0.6239137298E+00, 0.1688554040E+00]
+ζ = reshape(ζ, 1, length(ζ))
+D = [0.1543289673E+00, 0.5353281423E+00, 0.4446345422E+00]
+D = reshape(D, 1, length(D))
+
+Dn = STO_NG((ζ, D))
+bRnl = AtomicOrbitalsRadials(Pn, Dn, spec) 
+bYlm = RYlmBasis(Ylmdegree)
+
+# setup state
+ord = 1
+wf, spec, spec1p = BFwf_chain, spec, spec1p  = BFwf_lux(Nel, bRnl, bYlm, nuclei; totdeg = totdegree, ν = ord)
+
+ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ)
+p, = destructure(ps)
+length(p)
+
+K(wf, X::AbstractVector, ps, st) = -0.5 * laplacian(wf, X, ps, st)
+Vext(wf, X::AbstractVector, ps, st) = -sum(nuclei[i].charge/norm(nuclei[i].rr - X[j]) for i = 1:length(nuclei) for j in 1:length(X))
+function Vee(wf, X::AbstractVector, ps, st) 
+    nX = length(X)
+    if nX <=1 
+        return 0
+    else 
+        return sum(1/norm(X[i]-X[j]) for i = 1:length(X)-1 for j = i+1:length(X))
+    end
+end
+
+ham = SumH(K, Vext, Vee)
+sam = MHSampler(wf, Nel, nuclei, Δt = 1.0, burnin = 1000, nchains = 2000)
+
+opt_vmc = VMC(3000, 0.1, ACEpsi.vmc.adamW(); lr_dc = 100.0)
+wf, err_opt, ps = gd_GradientByVMC(opt_vmc, sam, ham, wf, ps, st)