Tucker format

Project.toml

@@ -9,6 +9,7 @@ BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
+Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 HyperDualNumbers = "50ceba7f-c3ee-5a84-a6e8-3ad40456ec97"
 IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"
@@ -20,13 +21,15 @@ LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
 ObjectPools = "658cac36-ff0f-48ad-967c-110375d98c9d"
 Octavian = "6fd5a793-0b7e-452c-907f-f8bfe9c57db4"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
+Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 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"
+Tullio = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]

benchmarks/benchmark_bflow_lux_Tucker.jl

@@ -0,0 +1,69 @@
+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 ACEpsi.TD: Tucker
+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 = 20
+Nel = 10
+X = randn(SVector{3, Float64}, Nel)
+Σ = rand(spins(), Nel)
+nuclei = [ Nuc(3 * rand(SVector{3, Float64}), 1.0) for _=1:10 ]
+
+# 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 = ACEpsi.BFwf_lux(Nel, bRnl, bYlm, nuclei, Tucker(5); totdeg = totdegree, ν = 2)
+ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ)
+out, st = BFwf_chain(X, ps, st)
+
+
+@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
+
+@profview let  BFwf_chain = BFwf_chain, X = X, ps =  ps, st = st
+   for i = 1:10_000
+      gradient(BFwf_chain, X, ps, st)
+   end
+end
+
+@profview let  BFwf_chain = BFwf_chain, X = X, ps =  ps, st = st
+   for i = 1:10_000
+      laplacian(BFwf_chain, X, ps, st)
+   end
+end

examples/3D/Be.jl

@@ -1,62 +1,74 @@
+using Distributed
+
+N_procs = 8
+
+if nprocs() == 1
+    addprocs(N_procs - 1, exeflags="--project=$(Base.active_project())")
+end
+
+@everywhere begin
+
 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 ACEpsi: BackflowPooling, BFwf_lux, setupBFState, Jastrow, displayspec
+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 Optimisers
 using Random
 using Printf
 using LinearAlgebra
 using BenchmarkTools
 using HyperDualNumbers: Hyper
-
-n1 = Rnldegree = 5
-Ylmdegree = 2
-totdegree = 5
+end 
+@everywhere begin
 Nel = 4
 X = randn(SVector{3, Float64}, Nel)
 Σ = [↑,↑,↓,↓]
 nuclei = [ Nuc(zeros(SVector{3, Float64}), Nel * 1.0)]
 ##
 
-# Defining AtomicOrbitalsBasis
-n2 = 1
+spec = [(n1 = 1, n2 = 1, l = 0), (n1 = 2, n2 = 1, l = 0), (n1 = 2, n2 = 1, l = 1), (n1 = 3, n2 = 1, l = 0)]
+n1 = 3
 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] 
+Ylmdegree = 2
+totdegree = 20
 ζ = 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)
+totdegree = [30,30,30]
+ν = [1,1,2]
+MaxIters = [100,100,200]
+_spec = [spec[1:3], spec, spec]
+_TD = [ACEpsi.Tucker(5),ACEpsi.Tucker(6),ACEpsi.Tucker(7)]
+wf_list, spec_list, spec1p_list, specAO_list, ps_list, st_list = wf_multilevel(Nel, Σ, nuclei, Dn, Pn, bYlm, _spec, totdegree, ν, _TD)
 
-ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ)
-p, = destructure(ps)
-length(p)
+ham = SumH(nuclei)
+sam = MHSampler(wf_list[1], Nel, nuclei, Δt = 0.5, burnin = 1000, nchains = 2000)
+opt_vmc = VMC_multilevel(MaxIters, 0.2, ACEpsi.vmc.adamW(); lr_dc = 50.0)
+end
+wf, err_opt, ps = gd_GradientByVMC_multilevel(opt_vmc, sam, ham, wf_list, ps_list, 
+                    st_list, spec_list, spec1p_list, specAO_list, batch_size = 500)
 
-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))
+# Eref = -14.667
+# HF   = -14.573
 
-ham = SumH(nuclei)
-sam = MHSampler(wf, Nel, nuclei, Δt = 0.5, burnin = 1000, nchains = 2000)
+# -14.12954986600700, var = 0.08884, 2p+js, ord = 1, size of basis = 5, 25 parameters
+# -14.60038627341732, var = 0.03319, 2p+js, ord = 2, size of basis = 70, 285 parameters
 
-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)
+# -14.493397693794488, var = 0.05635, 3s+js, ord = 1, size of basis = 6, 30 parameters
+# -14.627047427906787, var = 0.01082, 3s+js, ord = 2, size of basis = 99, 401 parameters
+# -14.631736047589962, var = 0.00883, 3s+js, ord = 2, size of basis = 99, 400 parameters
+
+
+# -14.368884109589212, var = 0.00935, 3s+js+1s*2, ord = 2, size of basis = 133, 539 parameters
+
+# -14.626306117451355, var = 0.01229, 3p+js, ord = 1, size of basis = 9, 43 parameters
+# -14.628261819128742, var = 0.01130, 3p+js, ord = 2, size of basis = 216, 871 parameters
 
-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
 

examples/3D/He.jl

@@ -0,0 +1,61 @@
+using Distributed
+
+N_procs = 10
+
+if nprocs() == 1
+    addprocs(N_procs - 1, exeflags="--project=$(Base.active_project())")
+end
+
+@everywhere begin
+
+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, 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
+using Random
+using Printf
+using LinearAlgebra
+using BenchmarkTools
+using HyperDualNumbers: Hyper
+
+end 
+@everywhere begin
+Nel = 2
+X = randn(SVector{3, Float64}, Nel)
+Σ = [↑,↓]
+nuclei = [ Nuc(zeros(SVector{3, Float64}), Nel * 1.0)]
+##
+
+spec = [(n1 = 1, n2 = 1, l = 0), (n1 = 2, n2 = 1, l = 0), (n1 = 2, n2 = 1, l = 1), (n1 = 3, n2 = 1, l = 0), (n1 = 3, n2 = 1, l = 1)]
+n1 = 3
+Pn = Polynomials4ML.legendre_basis(n1+1)
+Ylmdegree = 2
+totdegree = 20
+ζ = 10 * rand(length(spec))
+Dn = SlaterBasis(ζ)
+bYlm = RYlmBasis(Ylmdegree)
+
+totdegree = [30,30,30]
+ν = [1,1,2]
+MaxIters = [100,100,500]
+_spec = [spec[1:3], spec[1:4], spec]
+_TD = [ACEpsi.Tucker(5),ACEpsi.Tucker(6),ACEpsi.Tucker(7)]
+wf_list, spec_list, spec1p_list, specAO_list, ps_list, st_list = wf_multilevel(Nel, Σ, nuclei, Dn, Pn, bYlm, _spec, totdegree, ν, _TD)
+
+ham = SumH(nuclei)
+sam = MHSampler(wf_list[1], Nel, nuclei, Δt = 0.5, burnin = 1000, nchains = 2000)
+opt_vmc = VMC_multilevel(MaxIters, 0.5, ACEpsi.vmc.adamW(); lr_dc = 50.0)
+end
+wf, err_opt, ps = gd_GradientByVMC_multilevel(opt_vmc, sam, ham, wf_list, ps_list, 
+                    st_list, spec_list, spec1p_list, specAO_list, batch_size = 500)
+
+# Eref = -2.9037247 He
+# Eref = −7.47798 Li
+# 3s+js, ord = 2, -2.8990571700941272, 205 parameters He
+# 3p+js, ord = 2, -2.899325323029854, 440 parameters He