add benchmark for tucker bf

benchmarks/benchmark_bflow_lux_Tucker.jl

@@ -0,0 +1,67 @@
+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 = BFwf_lux(Nel, bRnl, bYlm, nuclei, Tucker(5); totdeg = totdegree, ν = 2)
+ps, st = setupBFState(MersenneTwister(1234), BFwf_chain, Σ)
+
+@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(X, ps, st)
+   end
+end

src/ACEpsi.jl

@@ -15,10 +15,11 @@ include("atomicorbitals/atomicorbitals.jl")
 include("jastrow.jl")
 include("tensordecomposition/TD.jl")
 include("bflow3d.jl")
-include("mbflow3d.jl")
-
 include("backflowpooling.jl")
 
+# experimental wfs
+include("experimental/mbflow3d.jl")
+
 # lux utils for bflow
 include("lux_utils.jl")