logistic and fix test

ACEsuit · Nov 26, 2023 · ae1a6d7 · ae1a6d7
1 parent bab1b27
commit ae1a6d7
Show file tree

Hide file tree

Showing 5 changed files with 111 additions and 101 deletions.
diff --git a/src/ACEpsi.jl b/src/ACEpsi.jl
@@ -17,12 +17,13 @@ include("tensordecomposition/TD.jl")
 include("bflow3d.jl")
 include("backflowpooling.jl")
 
-# experimental wfs
-include("experimental/mbflow3d.jl")
-
 # lux utils for bflow
 include("lux_utils.jl")
 
+# experimental wfs
+include("experimental/mbflow3d.jl")
+include("experimental/mmultilevel.jl")
+
 # vmc
 include("vmc/opt.jl")
 

diff --git a/src/experimental/mmultilevel.jl b/src/experimental/mmultilevel.jl
@@ -0,0 +1,104 @@
+using Printf
+using LinearAlgebra
+using Optimisers
+using Polynomials4ML
+using Random
+using ACEpsi: BackflowPooling, BFwf_lux, setupBFState, Jastrow, displayspec, mBFwf, mBFwf_sto
+using ACEpsi.AtomicOrbitals: _invmap
+using Plots
+
+
+function mwf_multilevel_sto(Nel::Int, Σ::Vector{Char}, nuclei::Vector{Nuc{T}}, Nbf::Vector{Int}, 
+                        Dn::STO_NG,
+                        Pn::OrthPolyBasis1D3T,  
+                        bYlm::Union{RYlmBasis, CYlmBasis, CRlmBasis},
+                        _spec::Vector{Vector{NamedTuple{(:n1, :n2, :l), Tuple{Int64, Int64, Int64}}}}, 
+                        totdegree::Vector{Int}, 
+                        ν::Vector{Int}
+                        ) where {T}
+    level = length(ν)
+    wf, spec, spec1p, ps, st = [], [], [], [], []
+    for i = 1:level
+        bRnl = AtomicOrbitalsRadials(Pn, Dn, _spec[i])
+        _wf, _spec1, _spec1p = mBFwf_sto(Nel, bRnl, bYlm, nuclei, Nbf[i]; totdeg = totdegree[i], ν = ν[i])
+        _ps, _st = setupBFState(MersenneTwister(1234), _wf, Σ)
+        push!(wf, _wf)
+        push!(spec, _spec1)
+        push!(spec1p, _spec1p)
+        push!(ps, _ps)
+        push!(st, _st)
+    end
+    return wf, spec, spec1p, _spec, ps, st
+end
+
+function mwf_multilevel(Nel::Int, Σ::Vector{Char}, nuclei::Vector{Nuc{T}}, Nbf::Vector{Int},  
+                        Dn::STO_NG,
+                        Pn::OrthPolyBasis1D3T,  
+                        bYlm::Union{RYlmBasis, CYlmBasis, CRlmBasis},
+                        _spec::Vector{Vector{NamedTuple{(:n1, :n2, :l), Tuple{Int64, Int64, Int64}}}}, 
+                        totdegree::Vector{Int}, 
+                        ν::Vector{Int}) where {T}
+    level = length(ν)
+    wf, spec, spec1p, ps, st = [], [], [], [], []
+    for i = 1:level
+        bRnl = AtomicOrbitalsRadials(Pn, Dn, _spec[i])
+        _wf, _spec1, _spec1p = mBFwf(Nel, bRnl, bYlm, nuclei, Nbf[i]; totdeg = totdegree[i], ν = ν[i])
+        _ps, _st = setupBFState(MersenneTwister(1234), _wf, Σ)
+        push!(wf, _wf)
+        push!(spec, _spec1)
+        push!(spec1p, _spec1p)
+        push!(ps, _ps)
+        push!(st, _st)
+    end
+    return wf, spec, spec1p, _spec, ps, st
+end
+
+# GaussianBasis
+function mwf_multilevel(Nel::Int, Σ::Vector{Char}, nuclei::Vector{Nuc{T}}, Nbf::Vector{Int},  
+                        Dn::GaussianBasis,
+                        Pn::OrthPolyBasis1D3T,  
+                        bYlm::Union{RYlmBasis, CYlmBasis, CRlmBasis},
+                        _spec::Vector{Vector{NamedTuple{(:n1, :n2, :l), Tuple{Int64, Int64, Int64}}}}, 
+                        totdegree::Vector{Int}, 
+                        ν::Vector{Int}) where {T}
+    level = length(ν)
+    wf, spec, spec1p, ps, st = [], [], [], [], []
+    for i = 1:level
+        ζ =  ones(Float64,length(_spec[i]))
+        Dn = GaussianBasis(ζ)
+        bRnl = AtomicOrbitalsRadials(Pn, Dn, _spec[i])
+        _wf, _spec1, _spec1p = mBFwf(Nel, bRnl, bYlm, nuclei, Nbf[i]; totdeg = totdegree[i], ν = ν[i])
+        _ps, _st = setupBFState(MersenneTwister(1234), _wf, Σ)
+        push!(wf, _wf)
+        push!(spec, _spec1)
+        push!(spec1p, _spec1p)
+        push!(ps, _ps)
+        push!(st, _st)
+    end
+    return wf, spec, spec1p, _spec, ps, st
+end
+
+# SlaterBasis
+function mwf_multilevel(Nel::Int, Σ::Vector{Char}, nuclei::Vector{Nuc{T}}, Nbf::Vector{Int},  
+                        Dn::SlaterBasis,
+                        Pn::OrthPolyBasis1D3T,  
+                        bYlm::Union{RYlmBasis, CYlmBasis, CRlmBasis},
+                        _spec::Vector{Vector{NamedTuple{(:n1, :n2, :l), Tuple{Int64, Int64, Int64}}}}, 
+                        totdegree::Vector{Int}, 
+                        ν::Vector{Int}) where {T}
+    level = length(ν)
+    wf, spec, spec1p, ps, st = [], [], [], [], []
+    for i = 1:level
+        ζ = ones(Float64,length(_spec[i]))
+        Dn = SlaterBasis(ζ)
+        bRnl = AtomicOrbitalsRadials(Pn, Dn, _spec[i])
+        _wf, _spec1, _spec1p = mBFwf(Nel, bRnl, bYlm, nuclei, Nbf[i]; totdeg = totdegree[i], ν = ν[i])
+        _ps, _st = setupBFState(MersenneTwister(1234), _wf, Σ)
+        push!(wf, _wf)
+        push!(spec, _spec1)
+        push!(spec1p, _spec1p)
+        push!(ps, _ps)
+        push!(st, _st)
+    end
+    return wf, spec, spec1p, _spec, ps, st
+end
diff --git a/src/vmc/multilevel.jl b/src/vmc/multilevel.jl
@@ -258,100 +258,4 @@ function wf_multilevel(Nel::Int, Σ::Vector{Char}, nuclei::Vector{Nuc{T}},
         push!(st, _st)
     end
     return wf, spec, spec1p, _spec, ps, st
-end
-
-
-function mwf_multilevel_sto(Nel::Int, Σ::Vector{Char}, nuclei::Vector{Nuc{T}}, Nbf::Vector{Int}, 
-                        Dn::STO_NG,
-                        Pn::OrthPolyBasis1D3T,  
-                        bYlm::Union{RYlmBasis, CYlmBasis, CRlmBasis},
-                        _spec::Vector{Vector{NamedTuple{(:n1, :n2, :l), Tuple{Int64, Int64, Int64}}}}, 
-                        totdegree::Vector{Int}, 
-                        ν::Vector{Int}
-                        ) where {T}
-    level = length(ν)
-    wf, spec, spec1p, ps, st = [], [], [], [], []
-    for i = 1:level
-        bRnl = AtomicOrbitalsRadials(Pn, Dn, _spec[i])
-        _wf, _spec1, _spec1p = mBFwf_sto(Nel, bRnl, bYlm, nuclei, Nbf[i]; totdeg = totdegree[i], ν = ν[i])
-        _ps, _st = setupBFState(MersenneTwister(1234), _wf, Σ)
-        push!(wf, _wf)
-        push!(spec, _spec1)
-        push!(spec1p, _spec1p)
-        push!(ps, _ps)
-        push!(st, _st)
-    end
-    return wf, spec, spec1p, _spec, ps, st
-end
-
-function mwf_multilevel(Nel::Int, Σ::Vector{Char}, nuclei::Vector{Nuc{T}}, Nbf::Vector{Int},  
-                        Dn::STO_NG,
-                        Pn::OrthPolyBasis1D3T,  
-                        bYlm::Union{RYlmBasis, CYlmBasis, CRlmBasis},
-                        _spec::Vector{Vector{NamedTuple{(:n1, :n2, :l), Tuple{Int64, Int64, Int64}}}}, 
-                        totdegree::Vector{Int}, 
-                        ν::Vector{Int}) where {T}
-    level = length(ν)
-    wf, spec, spec1p, ps, st = [], [], [], [], []
-    for i = 1:level
-        bRnl = AtomicOrbitalsRadials(Pn, Dn, _spec[i])
-        _wf, _spec1, _spec1p = mBFwf(Nel, bRnl, bYlm, nuclei, Nbf[i]; totdeg = totdegree[i], ν = ν[i])
-        _ps, _st = setupBFState(MersenneTwister(1234), _wf, Σ)
-        push!(wf, _wf)
-        push!(spec, _spec1)
-        push!(spec1p, _spec1p)
-        push!(ps, _ps)
-        push!(st, _st)
-    end
-    return wf, spec, spec1p, _spec, ps, st
-end
-
-# GaussianBasis
-function mwf_multilevel(Nel::Int, Σ::Vector{Char}, nuclei::Vector{Nuc{T}}, Nbf::Vector{Int},  
-                        Dn::GaussianBasis,
-                        Pn::OrthPolyBasis1D3T,  
-                        bYlm::Union{RYlmBasis, CYlmBasis, CRlmBasis},
-                        _spec::Vector{Vector{NamedTuple{(:n1, :n2, :l), Tuple{Int64, Int64, Int64}}}}, 
-                        totdegree::Vector{Int}, 
-                        ν::Vector{Int}) where {T}
-    level = length(ν)
-    wf, spec, spec1p, ps, st = [], [], [], [], []
-    for i = 1:level
-        ζ =  ones(Float64,length(_spec[i]))
-        Dn = GaussianBasis(ζ)
-        bRnl = AtomicOrbitalsRadials(Pn, Dn, _spec[i])
-        _wf, _spec1, _spec1p = mBFwf(Nel, bRnl, bYlm, nuclei, Nbf[i]; totdeg = totdegree[i], ν = ν[i])
-        _ps, _st = setupBFState(MersenneTwister(1234), _wf, Σ)
-        push!(wf, _wf)
-        push!(spec, _spec1)
-        push!(spec1p, _spec1p)
-        push!(ps, _ps)
-        push!(st, _st)
-    end
-    return wf, spec, spec1p, _spec, ps, st
-end
-
-# SlaterBasis
-function mwf_multilevel(Nel::Int, Σ::Vector{Char}, nuclei::Vector{Nuc{T}}, Nbf::Vector{Int},  
-                        Dn::SlaterBasis,
-                        Pn::OrthPolyBasis1D3T,  
-                        bYlm::Union{RYlmBasis, CYlmBasis, CRlmBasis},
-                        _spec::Vector{Vector{NamedTuple{(:n1, :n2, :l), Tuple{Int64, Int64, Int64}}}}, 
-                        totdegree::Vector{Int}, 
-                        ν::Vector{Int}) where {T}
-    level = length(ν)
-    wf, spec, spec1p, ps, st = [], [], [], [], []
-    for i = 1:level
-        ζ = ones(Float64,length(_spec[i]))
-        Dn = SlaterBasis(ζ)
-        bRnl = AtomicOrbitalsRadials(Pn, Dn, _spec[i])
-        _wf, _spec1, _spec1p = mBFwf(Nel, bRnl, bYlm, nuclei, Nbf[i]; totdeg = totdegree[i], ν = ν[i])
-        _ps, _st = setupBFState(MersenneTwister(1234), _wf, Σ)
-        push!(wf, _wf)
-        push!(spec, _spec1)
-        push!(spec1p, _spec1p)
-        push!(ps, _ps)
-        push!(st, _st)
-    end
-    return wf, spec, spec1p, _spec, ps, st
-end
+end
diff --git a/src/vmc/vmc.jl b/src/vmc/vmc.jl
@@ -3,6 +3,7 @@ using Printf
 using LinearAlgebra
 using Optimisers
 using Plots
+
 mutable struct VMC
     tol::Float64
     MaxIter::Int

diff --git a/test/test_atorbbasis.jl b/test/test_atorbbasis.jl
@@ -84,7 +84,7 @@ for I = 1:Nnuc
                a2 = Rnl[i, iR] * Ylm[i, iY]
             end
             # println("(i=$i, σ=$s, I=$I, n=$(nlm.n), l=$(nlm.l), m=$(nlm.m)) -> ", abs(a1 - a2))
-            print_tf(@test a1 ≈ a2)
+            print_tf(@test norm(a1 - a2) < 1e-12)
          end
       end
    end