Merge pull request #22 from ACEsuit/test_noparallel

3D BackFlowLux
ACEsuit · Oct 5, 2023 · 11661dc · 11661dc
2 parents 312bf26 + 7f53212
commit 11661dc
Show file tree

Hide file tree

Showing 40 changed files with 2,780 additions and 136 deletions.
diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
@@ -18,8 +18,8 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.0'
           - '1.8'
+          - '1.9'
           - 'nightly'
         os:
           - ubuntu-latest

diff --git a/Project.toml b/Project.toml
@@ -5,25 +5,37 @@ 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"
+IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"
 JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+LinearMaps = "7a12625a-238d-50fd-b39a-03d52299707e"
+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"

diff --git a/benchmarks/benchmark_bflow_lux.jl b/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
diff --git a/examples/3D/Be.jl b/examples/3D/Be.jl
@@ -0,0 +1,62 @@
+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(nuclei)
+sam = MHSampler(wf, Nel, nuclei, Δ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
+
diff --git a/examples/3D/He_multi.jl b/examples/3D/He_multi.jl
@@ -0,0 +1,45 @@
+using ACEpsi, Polynomials4ML, StaticArrays, Test 
+using Polynomials4ML: natural_indices, degree, SparseProduct
+using ACEpsi.AtomicOrbitals: Nuc, make_nlms_spec, evaluate
+using ACEpsi.vmc: gradient, laplacian, grad_params, EmbeddingW!, _invmap, VMC_multilevel, wf_multilevel, VMC, gd_GradientByVMC, gd_GradientByVMC_multilevel, AdamW, SR, SumH, MHSampler
+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
+
+
+# Define He model 
+Nel = 2
+X = randn(SVector{3, Float64}, Nel)
+Σ = [↑,↓]
+nuclei = [ Nuc(zeros(SVector{3, Float64}), 2.0)]
+
+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(nuclei)
+
+# Defining Multilevel
+Rnldegree = [4, 6, 6, 7]
+Ylmdegree = [2, 2, 3, 4]
+totdegree = [2, 3, 3, 4]
+n2 = [1, 1, 2, 2]
+ν = [1, 1, 2, 2]
+MaxIters = [3, 3, 3, 3]
+##
+
+# 
+wf_list, spec_list, spec1p_list, specAO_list, ps_list, st_list = wf_multilevel(Nel, Σ, nuclei, Rnldegree, Ylmdegree, totdegree, n2, ν)
+
+sam = MHSampler(wf_list[1], Nel, nuclei, Δt = 0.5, burnin = 1, nchains = 20)
+opt_vmc = VMC_multilevel(MaxIters, 0.0015, SR(1e-4, 0.015), lr_dc = 50.0)
+
+wf, err_opt, ps = gd_GradientByVMC_multilevel(opt_vmc, sam, ham, wf_list, ps_list, st_list, spec_list, spec1p_list, specAO_list)
+
diff --git a/examples/3D/ccpvdz_H10.jl b/examples/3D/ccpvdz_H10.jl
@@ -0,0 +1,80 @@
+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
+using Random
+using Printf
+using LinearAlgebra
+using BenchmarkTools
+using HyperDualNumbers: Hyper
+
+n1 = Rnldegree = 2
+Ylmdegree = 1
+totdegree = 30
+Nel = 10
+X = randn(SVector{3, Float64}, Nel)
+Σ = [↑,↑,↑,↑,↑,↓,↓,↓,↓,↓]
+spacing = 1.0
+nuclei = [Nuc(SVector(0.0,0.0,(i-1/2-Nel/2) * spacing), 1.0) for i = 1:Nel]
+Pn = Polynomials4ML.legendre_basis(n1)
+spec = [(n1 = 1, n2 = 1, l = 0), (n1 = 1, n2 = 2, l = 0), (n1 = 2, n2 = 1, l = 1)]
+
+# Ref: http://www.grant-hill.group.shef.ac.uk/ccrepo/hydrogen/hbasis.php
+# (4s,1p) -> [2s,1p]
+# H    S
+#      1.301000E+01           1.968500E-02           0.000000E+00
+#      1.962000E+00           1.379770E-01           0.000000E+00
+#      4.446000E-01           4.781480E-01           0.000000E+00
+#      1.220000E-01           5.012400E-01           1.000000E+00
+# H    P
+#      7.270000E-01           1.0000000
+
+ζ = [[1.301000E+01, 1.962000E+00, 4.446000E-01, 1.220000E-01], [1.220000E-01], [7.270000E-01]]
+D = [[1.968500E-02, 1.379770E-01, 4.781480E-01, 5.012400E-01], [1.0000000], [1.0000000]]
+D[1] = [(2 * ζ[1][i]/pi)^(3/4) * D[1][i] for i = 1:length(ζ[1])] * sqrt(2) * 2 * sqrt(pi)
+D[2] = [(2 * ζ[2][i]/pi)^(3/4) * D[2][i] for i = 1:length(ζ[2])] * sqrt(2) * 2 * sqrt(pi)
+
+Dn = STO_NG((ζ, D))
+bRnl = AtomicOrbitalsRadials(Pn, Dn, spec) 
+bYlm = RYlmBasis(Ylmdegree)
+
+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)
+wf(X, ps, st)
+@profview begin for i = 1:100000 wf(X, ps, st) end end
+
+ham = SumH(nuclei)
+sam = MHSampler(wf, Nel, nuclei, Δt = 0.5, burnin = 1000, nchains = 2000)
+
+#using BenchmarkTools # N = 10, nuc = 10  # 50, 6325, 408425
+#@btime $wf($X, $ps, $st)               # ord = 1: 22.764 μs (17 allocations: 19.67 KiB)
+                                        # ord = 2: 119.143 μs (17 allocations: 19.67 KiB)
+                                        # ord = 3: 9.800 ms (17 allocations: 19.67 KiB)
+
+#@btime $gradient($wf, $X, $ps, $st)    # ord = 1: 81.785 μs (165 allocations: 153.27 KiB)
+                                        # ord = 2: 742.233 μs (167 allocations: 1.11 MiB)
+                                        # ord = 3: 60.314 ms (167 allocations: 62.47 MiB)
+
+#@btime $grad_params($wf, $X, $ps, $st) # ord = 1: 74.807 μs (142 allocations: 94.42 KiB)
+                                        # ord = 2: 732.846 μs (144 allocations: 1.05 MiB)
+                                        # ord = 3: 59.758 ms (144 allocations: 62.41 MiB)
+
+#@btime $laplacian($wf, $X, $ps, $st)   # ord = 1: 1.215 ms (611 allocations: 2.20 MiB)
+                                        # ord = 2: 25.875 ms (611 allocations: 2.20 MiB)
+                                        # ord = 3: 2.648 s (611 allocations: 2.20 MiB)
+
+opt_vmc = VMC(5000, 0.015, ACEpsi.vmc.adamW(); lr_dc = 300.0)
+#wf, err_opt, ps = gd_GradientByVMC(opt_vmc, sam, ham, wf, ps, st)
+
+## MRCI+Q: -23.5092
+## UHF:    -23.2997
diff --git a/examples/3D/sto-6g_H10.jl b/examples/3D/sto-6g_H10.jl
@@ -0,0 +1,94 @@
+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
+using Random
+using Printf
+using LinearAlgebra
+using BenchmarkTools
+using HyperDualNumbers: Hyper
+
+n1 = Rnldegree = 1
+Ylmdegree = 0
+totdegree = 20
+Nel = 10
+X = randn(SVector{3, Float64}, Nel)
+Σ = [↑,↑,↑,↑,↑,↓,↓,↓,↓,↓]
+spacing = 1.0
+nuclei = [Nuc(SVector(0.0,0.0,(i-1/2-Nel/2) * spacing), 1.0) for i = 1:Nel]
+Pn = Polynomials4ML.legendre_basis(n1+1)
+spec = [(n1 = 1, n2 = 1, l = 0)] 
+
+# Ref: https://link.springer.com/book/10.1007/978-90-481-3862-3: P235
+# STO: 0.7790 * e^(-1.24 * r)
+# ϕ_1s(1, r) = \sum_(k = 1)^K d_1s,k g_1s(α_1k, r)
+# g_1s(α, r) = (2α/π)^(3/4) * exp(-αr^2): α ∼ ζ, g ∼ D
+# sto-3g: Ref: https://www.basissetexchange.org/
+# BASIS SET: (3s) -> [1s]
+# H    S
+#      0.3425250914E+01       0.1543289673E+00
+#      0.6239137298E+00       0.5353281423E+00
+#      0.1688554040E+00       0.4446345422E+00
+
+# ζ = [[0.3425250914E+01, 0.6239137298E+00, 0.1688554040E+00]]
+# D = [[0.1543289673E+00, 0.5353281423E+00, 0.4446345422E+00]]
+# D[1] = [(2 * ζ[1][i]/pi)^(3/4) * D[1][i] for i = 1:length(ζ[1])]
+# P_0(x) = 1/sqrt(2)
+# Y_0(x) = 1/(2*sqrt(pi))
+
+# sto-6g: Ref: https://www.basissetexchange.org/
+# BASIS SET: (6s) -> [1s]
+# H    S
+#      0.3552322122E+02       0.9163596281E-02
+#      0.6513143725E+01       0.4936149294E-01
+#      0.1822142904E+01       0.1685383049E+00
+#      0.6259552659E+00       0.3705627997E+00
+#      0.2430767471E+00       0.4164915298E+00
+#      0.1001124280E+00       0.1303340841E+00
+
+ζ = [[0.3552322122E+02, 0.6513143725E+01, 0.1822142904E+01,0.6259552659E+00, 0.2430767471E+00, 0.1001124280E+00]]
+D = [[0.9163596281E-02, 0.4936149294E-01,0.1685383049E+00,0.3705627997E+00,  0.4164915298E+00, 0.1303340841E+00]]
+D[1] = [(2 * ζ[1][i]/pi)^(3/4) * D[1][i] for i = 1:length(ζ[1])] * sqrt(2) * 2 * sqrt(pi)
+
+Dn = STO_NG((ζ, D))
+bRnl = AtomicOrbitalsRadials(Pn, Dn, spec) 
+bYlm = RYlmBasis(Ylmdegree)
+
+ord = 3
+wf, spec, spec1p = BFwf_chain, spec, spec1p  = BFwf_lux(Nel, bRnl, bYlm, nuclei; totdeg = totdegree, ν = ord)
+
+ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ) # ps.hidden1.W: Nels * basis
+p, = destructure(ps)
+length(p)
+
+ham = SumH(nuclei)
+sam = MHSampler(wf, Nel, nuclei, Δt = 0.5, burnin = 1000, nchains = 2000)
+
+#using BenchmarkTools # N = 10, nuc = 10  # 10, 265, 3685
+#@btime $wf($X, $ps, $st)               # ord = 1: 17.197 μs (17 allocations: 6.86 KiB)
+                                        # ord = 2: 23.086 μs (17 allocations: 6.86 KiB)
+                                        # ord = 3: 78.264 μs (17 allocations: 6.86 KiB)
+
+#@btime $gradient($wf, $X, $ps, $st)    # ord = 1: 56.969 μs (165 allocations: 100.73 KiB)
+                                        # ord = 2: 75.596 μs (167 allocations: 143.98 KiB)
+                                        # ord = 3: 889.464 μs (167 allocations: 681.77 KiB)
+
+#@btime $grad_params($wf, $X, $ps, $st) # ord = 1: 48.483 μs (142 allocations: 41.89 KiB)
+                                        # ord = 2: 67.085 μs (144 allocations: 81.73 KiB)
+                                        # ord = 3: 887.024 μs (144 allocations: 616.11 KiB)
+
+#@btime $laplacian($wf, $X, $ps, $st)   # ord = 1: 738.616 μs (581 allocations: 720.28 KiB)
+                                        # ord = 2: 1.747 ms (581 allocations: 720.28 KiB)
+                                        # ord = 3: 20.752 ms (581 allocations: 720.28 KiB)
+
+opt_vmc = VMC(5000, 0.015, ACEpsi.vmc.adamW(); lr_dc = 100.0)
+wf, err_opt, ps = gd_GradientByVMC(opt_vmc, sam, ham, wf, ps, st)
+
+## FCI: -23.1140: ord = 2: -23.3829
+## UHF: -23.0414: ord = 1: -23.0432