add tensor

Author: DexuanZhou

Project.toml

@@ -26,6 +26,7 @@ ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+Tullio = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc"
 UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 

example.jl

@@ -5,17 +5,20 @@ using ACEpsi
 method = 1
 mol = ACEpsi.molecules.Be
 mol_name = "Be"
-setup(mol, mol_name, method)
+TD = SCPMultipleW(14) #No_Decomposition() 
+worldsize = N_procs = 8
+setup(mol, mol_name, method, TD, worldsize)
 
-x0, optimizer, model_list, ps_list, st_list, spec_list, spec1p_list, totdeg_list, ν_list = load_setup(mol_name);
+x0, optimizer, model_list, ps_list, st_list, spec_list, spec1p_list, totdeg_list, ν_list = load_setup(mol_name, TD);
 solver = (SPRINGSolver(), SketchSolver(800, 50, 50, 1.4), SVDSolver(800, 50, 50, 1.4))
 optimizer.sr_method = solver[method]
 string = method == 1 ? "SPRING" :
              method == 2 ? "SKETCH" :
              method == 3 ? "WSSR"   : error("Invalid method")
-optimizer.res_path = "$mol_name/$string/"
+clean_TD = replace("$(TD)", r"[^A-Za-z0-9]" => "")
+optimizer.res_path = "$mol_name/$string/$clean_TD/"
+
 
-N_procs = 8
 using Pkg
 using Distributed
 Pkg.activate(Base.current_project())

src/ACEpsi.jl

@@ -6,6 +6,7 @@ include("molecule/molecule.jl")
 include("molecule/molecules.jl")
 include("molecule/hamiltonian.jl")
 include("model/layers.jl")
+include("model/tensorlayers.jl")
 include("model/spec.jl")
 include("model/wavefunction.jl")
 include("model/multilevel.jl")

src/model/layers.jl

@@ -54,9 +54,16 @@ function evaluate(jl::JastrowLayer, X::Vector{SVector{3, TX}}, Σ::Vector{Char})
 
         γ += -c_ij / (1 + dist)
     end
-    return γ
+    return exp(γ)
 end
 
+struct BackflowPoolingLayer_TD{TT}<: AbstractLuxLayer
+    spec::Vector{TT}
+    Σ::Vector{Char}
+end
+
+(l::BackflowPoolingLayer_TD)(x, ps, st) = evaluate(l, x, l.Σ), ps, st
+
 function evaluate(l::BackflowPoolingLayer, x, Σ::Vector{Char})
     T = promote_type(eltype(x[1]))
     Nnlm = length(l.spec)
@@ -93,6 +100,43 @@ function evaluate(l::BackflowPoolingLayer, x, Σ::Vector{Char})
     return A  # shape: (Nel, 3 * Nnlm)
 end
 
+function evaluate(l::BackflowPoolingLayer_TD, x, Σ::Vector{Char})
+    T = promote_type(eltype(x[1]))
+    Nnlm = length(l.spec)
+    Nel = length(Σ)
+
+    @assert spin2idx(↑) == 1
+    @assert spin2idx(↓) == 2
+    @assert spin2idx(∅) == 3
+
+    A = zeros(T, Nel, 3, Nnlm)       # (Nel, 3 channel, Nnlm)
+    Aall = zeros(T, 2, Nnlm)         # (spin channel ∈ {↑, ↓}, Nnlm)
+
+    @inbounds begin
+        for k = 1:Nnlm
+            @simd ivdep for i = 1:Nel
+                iσ = spin2idx(Σ[i])
+                if iσ ≤ 2
+                    Aall[iσ, k] += x[i, k]
+                end
+                A[i, 3, k] = x[i, k]
+            end
+        end
+
+        for k = 1:Nnlm
+            @simd ivdep for iσ = 1:2
+                σ = idx2spin(iσ)
+                for i = 1:Nel
+                    A[i, iσ, k] = Aall[iσ, k] - (Σ[i] == σ ? x[i, k] : zero(T))
+                end
+            end
+        end
+    end
+
+    return A  # shape: (Nel, 3, Nnlm)
+end
+
+
 function ChainRulesCore.rrule(::typeof(evaluate), l::Diff_layer{Nnuc}, X::Vector{SVector{3, TX}}, ps::NamedTuple, st::NamedTuple) where {Nnuc, TX}
    val = ntuple(i -> _getdiff(X, l.nuc[i].rr), Val(Nnuc))
    function pb(dA)
@@ -119,6 +163,14 @@ function ChainRulesCore.rrule(::typeof(evaluate), pooling::BackflowPoolingLayer,
     return A, pb
 end 
 
+function ChainRulesCore.rrule(::typeof(evaluate), pooling::BackflowPoolingLayer_TD, x, Σ::Vector{Char}) 
+    A = evaluate(pooling, x, pooling.Σ)
+    function pb(∂A)
+        return NoTangent(), NoTangent(), _pullback_evaluate(∂A, pooling, x, Σ), NoTangent()
+    end
+    return A, pb
+end 
+
 # helper function
 
 function _getdiff(X::AbstractArray{SVector{3, T}}, d::SVector{3, TT}) where {T, TT}
@@ -156,6 +208,35 @@ function _pullback_evaluate(∂A, l::BackflowPoolingLayer, x, Σ::Vector{Char})
     return ∂x
 end
 
+function _pullback_evaluate(∂A, l::BackflowPoolingLayer_TD, x, Σ::Vector{Char})
+    TA = eltype(x[1])
+    Nel = length(Σ)
+    Nnlm = length(l.spec)
+
+    ∂x = zeros(TA, Nel, Nnlm)
+
+    @inbounds begin
+        for k = 1:Nnlm
+            for i = 1:Nel
+                σi = Σ[i]
+                iσ = spin2idx(σi)
+
+                ∂x[i, k] += ∂A[i, 3, k]
+
+                if iσ ≤ 2
+                    for j = 1:Nel
+                        if j != i && Σ[j] == σi
+                            ∂x[i, k] += ∂A[j, iσ, k]
+                        end
+                    end
+                end
+            end
+        end
+    end
+
+    return ∂x
+end
+
 
 function LuxCore.initialparameters(rng::AbstractRNG, d::Dense)
     weight = if d.init_weight === nothing

src/model/multilevel.jl

@@ -1,12 +1,35 @@
 export transfer_weights!, model_generator
 
 function transfer_weights!(ps1, ps2, spec1, spec2, spec1p1, spec1p2)
-    readable_spec1 = displayspec(spec1, spec1p1)
-    readable_spec2 = displayspec(spec2, spec1p2)
+    readable_spec1 = displayspec(spec1, spec1p1, ps1)
+    readable_spec2 = displayspec(spec2, spec1p2, ps2)
     _map  = _invmap(readable_spec2)
-    ps2.branch.bf.linear.weight .= 0.0
-    for (idx, t) in enumerate(readable_spec1)
-        ps2.branch.bf.linear.weight[:, _map[t]] = ps1.branch.bf.linear.weight[:, idx]
+    if :TK in keys(ps2.branch.bf)
+        _map2  = _invmap(spec1p2)
+        ps2.branch.bf.TK.W .= 0.0
+        W = ps1.branch.bf.TK.W
+        if length(size(W)) == 3
+            for (idx, t) in enumerate(spec1p1)
+                ps2.branch.bf.TK.W[:,1:size(W)[2], _map2[t]] .= W[:, 1:size(W)[2], idx]
+            end
+        elseif length(size(W)) == 4
+            for (idx, t) in enumerate(spec1p1)
+                ps2.branch.bf.TK.W[:, :, 1:size(W)[3], _map2[t]] .= W[:, :, 1:size(W)[3], idx]
+            end
+        end
+    end
+    if :linear in keys(ps2.branch.bf)
+        ps2.branch.bf.linear.weight .= 0.0
+        for (idx, t) in enumerate(readable_spec1)
+            ps2.branch.bf.linear.weight[:, _map[t]] = ps1.branch.bf.linear.weight[:, idx]
+        end
+    else
+        for i in keys(ps2.branch.bf.bAA)
+            ps2.branch.bf.bAA[i].layer_2.weight .= 0.0
+            for (idx, t) in enumerate(readable_spec1)
+                ps2.branch.bf.bAA[i].layer_2.weight[:, _map[t]] = ps1.branch.bf.bAA[i].layer_2.weight[:, idx]
+            end
+        end
     end
     return ps2
 end
@@ -19,18 +42,18 @@ function transfer_weights_idx!(ps1, ps2, spec1, spec2, spec1p1, spec1p2)
     return index
 end
 
-function model_generator(mol, basis_set, totdeg, ν; ratio = 0.5, multilevel = true, filename = "basis.json")
+function model_generator(mol, basis_set, totdeg, ν; TD = No_Decomposition(), ratio = 0.5, multilevel = true, filename = "basis.json")
     if multilevel
-        totdeg_list, ν_list = build_totdeglevels(mol, basis_set, totdeg, ν; ratio = ratio)
-        results = [build_wavefunction(mol, basis_set, totdeg_list[i], ν_list[i]; filename = filename) for i = 1:length(totdeg_list)]
+        totdeg_list, ν_list = build_totdeglevels(mol, basis_set, totdeg, ν, TD; ratio = ratio)
+        results = [build_wavefunction(mol, basis_set, totdeg_list[i], ν_list[i], TD; filename = filename) for i = 1:length(totdeg_list)]
         model_list = getindex.(results, 1)
         ps_list    = getindex.(results, 2)
         st_list    = getindex.(results, 3)
         spec_list = getindex.(results, 4)
         spec1p_list = getindex.(results, 5)
         return model_list, ps_list, st_list, spec_list, spec1p_list, totdeg_list, ν_list
     else
-        model, ps, st, spec, spec1p = build_wavefunction(mol, basis_set, totdeg, ν; filename = fieldname)
+        model, ps, st, spec, spec1p = build_wavefunction(mol, basis_set, totdeg, ν, TD; filename = fieldname)
         return [model], [ps], [st], [spec], [spec1p], [totdeg], [ν]
     end
 end

src/model/spec.jl

@@ -54,10 +54,24 @@ function parseTotdegToInt(
     return _totdegn
 end
 
-function displayspec(spec, spec1p)
-    nicespec = []
-    for k = 1:length(spec)
-        push!(nicespec, [spec1p[spec[k][j]] for j = 1:length(spec[k])])
+function displayspec(spec, spec1p, ps)
+    if :TK ∉ keys(ps.branch.bf)
+        nicespec = []
+        for k = 1:length(spec)
+            push!(nicespec, [spec1p[spec[k][j]] for j = 1:length(spec[k])])
+        end
+    else
+        P = 0
+        if length(size(ps.branch.bf.TK.W)) == 4
+            P = size(ps.branch.bf.TK.W)[3]
+        elseif length(size(ps.branch.bf.TK.W)) == 3
+            P = size(ps.branch.bf.TK.W)[2]
+        end
+        spec1p = get_spec1p(P)
+        nicespec = []
+        for k = 1:length(spec)
+            push!(nicespec, [spec1p[spec[k][j]] for j = 1:length(spec[k])])
+        end
     end
     return nicespec
 end
@@ -95,6 +109,18 @@ function get_spec1p(basis::Vector{TS}; spin = false) where {TS}
     return spec[:]
 end
 
+function get_spec1p(P)  
+    spec = Array{Any}(undef, (3, P))
+ 
+    for k = 1:P
+        for (is, s) in enumerate(extspins())
+            spec[is, k] = (s=s, P = k)
+        end
+    end
+ 
+   return spec[:]
+end
+
 function _invmap(a)
     inva = Dict{eltype(a), Int}()
     for i = 1:length(a) 
@@ -145,7 +171,7 @@ function sample_evenly(arr::AbstractVector; N = 10)
 end
 
 
-function build_totdeglevels(mol, basis_set, totdeg, ν; ratio = 0.5, max_level::Union{Nothing, Int} = nothing)
+function build_totdeglevels(mol, basis_set, totdeg, ν, TD::No_Decomposition; ratio = 0.5, max_level::Union{Nothing, Int} = nothing)
     _, orbital = auto_load_basis(mol, basis_set; return_spec = true)
     n_atom = length(orbital)
 
@@ -203,3 +229,45 @@ function build_totdeglevels(mol, basis_set, totdeg, ν; ratio = 0.5, max_level::
 
     return totdeglevels, νlevels
 end
+
+function build_totdeglevels(mol, basis_set, totdeg, ν, TD; ratio = 0.5, max_level::Union{Nothing, Int} = nothing)
+    maxdim = length(totdeg)
+    levels = Vector{Vector{Int}}()
+
+    d = mol.Nel
+    deg_split = floor(Int, totdeg[1] * ratio)
+    for x = d:deg_split
+        push!(levels, [x])
+    end
+
+    function extend_levels(levels, curdim)
+        result = Vector{Vector{Int}}()
+        for v in levels
+            if length(v) == curdim &&
+               all(v[i] ≥ floor(Int, totdeg[i] * ratio) for i in 1:curdim)
+                for j = mol.Nel:floor(Int, totdeg[curdim+1] * ratio)
+                    push!(result, vcat(v, j))
+                end
+            end
+        end
+        return result
+    end
+
+    for curdim = 1:maxdim-1
+        new_levels = extend_levels(levels, curdim)
+        append!(levels, new_levels)
+    end
+
+    last_diag = [floor(Int, totdeg[i] * ratio) for i in 1:maxdim]
+    while all(last_diag[i] ≤ totdeg[i] for i in 1:maxdim)
+        push!(levels, copy(last_diag))
+        for i in 1:maxdim
+            last_diag[i] += 1
+        end
+    end
+
+    νlevels = [length(v) for v in levels]
+    return levels, νlevels
+end
+
+