Skip to content

Commit

Permalink
Comment out those (L1,L2) block related codes
Browse files Browse the repository at this point in the history
  • Loading branch information
@zhanglw0521
zhanglw0521 committed Oct 21, 2023
1 parent 479a65c commit bf3a523
Show file tree
Hide file tree
Showing 2 changed files with 241 additions and 241 deletions.
374 changes: 187 additions & 187 deletions src/builder.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -305,191 +305,191 @@ equivariant_SYY_model(totdeg::Int64, ν::Int64, radial::Radial_basis, L::Int64;
equivariant_SYY_model(nn::Vector{Int64}, ll::Vector{Int64}, radial::Radial_basis, L::Int64; categories=[], _get_cat = _get_cat_default, d=3, group="O3") =
equivariant_SYY_model(_close(nn, ll; filter = RPE_filter_long(L)), radial, L; categories, d, group)

## TODO: The following should eventually go into ACEhamiltonians.jl rather than this package

# mapping from long vector to spherical matrix
using RepLieGroups.O3: ClebschGordan
cg = ClebschGordan(ComplexF64)
function cgmatrix(L1, L2)
cgm = zeros((2L1+1) * (2L2+1),(L1+L2+1)^2)
for (i,(p,q)) in enumerate(collect(Iterators.product(-L1:L1, L2:-1:-L2)))
ν = p+q
for l = 0:L1+L2
if abs(ν)<=l
position = ν + l + 1
cgm[i, l^2+position] = (-1)^q * cg(L1, p, L2, q, l, ν)
# cgm[i, l^2+position] = cg(L1,p,L2,q,l,ν)
# cgm[i, l^2+position] = (-1)^q * sqrt( (2L1+1) * (2L2+1) ) / 2 / sqrt(π * (2l+1)) * cg(L1,0,L2,0,l,0) * cg(L1,p,L2,q,l,ν)
end
end
end
return sparse(cgm)
end

# This constructor builds a lux chain that maps a configuration to a LONG B^L vector ([B^0, B^1, ... B^L]),
# and then all equivariant basis.
# What can be adjusted in its input are: (1) total polynomial degree; (2) correlation order; (3) largest L
# (4) weight of the order of spherical harmonics; (5) specified radial basis
function equivariant_luxchain_constructor(totdeg, ν, L; wL = 1, Rn = legendre_basis(totdeg))

filter = RPE_filter_long(L)
cgen = Rot3DCoeffs_long(L)

Ylm = CYlmBasis(totdeg)

spec1p = make_nlms_spec(simple_radial_basis(Rn), Ylm; totaldegree = totdeg, admissible = (br, by) -> br.n + wL * by.l <= totdeg)
spec1p = sort(spec1p, by = (x -> x.n + x.l * wL))
spec1pidx = getspec1idx(spec1p, simple_radial_basis(Rn).Radialspec, Ylm)

# define sparse for n-correlations
tup2b = vv -> [ spec1p[v] for v in vv[vv .> 0] ]
default_admissible = bb -> length(bb) == 0 || sum(b.n for b in bb) + wL * sum(b.l for b in bb) <= totdeg

specAA = gensparse(; NU = ν, tup2b = tup2b, filter = filter,
admissible = default_admissible,
minvv = fill(0, ν),
maxvv = fill(length(spec1p), ν),
ordered = true)

spec = [ vv[vv .> 0] for vv in specAA if !(isempty(vv[vv .> 0]))]
# map back to nlm
spec_nlm = getspecnlm(spec1p, spec)

C = _rpi_A2B_matrix(cgen, spec_nlm)

# acemodel with lux layers
bA = P4ML.PooledSparseProduct(spec1pidx)
bAA = P4ML.SparseSymmProd(spec)

# wrapping into lux layers
l_Rn = P4ML.lux(Rn)
l_Ylm = P4ML.lux(Ylm)
l_bA = P4ML.lux(bA)
l_bAA = P4ML.lux(bAA)

# formming model with Lux Chain
_norm(x) = norm.(x)

l_xnx = Lux.Parallel(nothing; normx = WrappedFunction(_norm), x = WrappedFunction(identity))
l_embed = Lux.Parallel(nothing; Rn = l_Rn, Ylm = l_Ylm)

l1l2set = [(l1,l2) for l1 = 0:L for l2 = 0:L-l1]
cgmat_set = [cgmatrix(l1,l2) for (l1,l2) in l1l2set ]

_block(x,cgmat,l1,l2) = reshape.(Ref(cgmat) .* [ x[j][1:size(cgmat)[2]] for j = 1:length(x) ], 2l1+1, 2l2+1)
_dropzero(x) = x[findall(X -> X>1e-12, norm.(x))]

l_seperate = Lux.Parallel(nothing, [WrappedFunction(x -> _dropzero(_block(x,cgmat_set[i],l1l2set[i][1],l1l2set[i][2]))) for i = 1:length(cgmat_set)]... )

# C - A2Bmap
luxchain = Chain(xnx = l_xnx, embed = l_embed, A = l_bA , AA = l_bAA, BB = WrappedFunction(x -> C * x), blocks = l_seperate)#, rAA = WrappedFunction(ComplexF64))
ps, st = Lux.setup(MersenneTwister(1234), luxchain)
# ## TODO: The following should eventually go into ACEhamiltonians.jl rather than this package

# # mapping from long vector to spherical matrix
# using RepLieGroups.O3: ClebschGordan
# cg = ClebschGordan(ComplexF64)
# function cgmatrix(L1, L2)
# cgm = zeros((2L1+1) * (2L2+1),(L1+L2+1)^2)
# for (i,(p,q)) in enumerate(collect(Iterators.product(-L1:L1, L2:-1:-L2)))
# ν = p+q
# for l = 0:L1+L2
# if abs(ν)<=l
# position = ν + l + 1
# cgm[i, l^2+position] = (-1)^q * cg(L1, p, L2, q, l, ν)
# # cgm[i, l^2+position] = cg(L1,p,L2,q,l,ν)
# # cgm[i, l^2+position] = (-1)^q * sqrt( (2L1+1) * (2L2+1) ) / 2 / sqrt(π * (2l+1)) * cg(L1,0,L2,0,l,0) * cg(L1,p,L2,q,l,ν)
# end
# end
# end
# return sparse(cgm)
# end

# # This constructor builds a lux chain that maps a configuration to a LONG B^L vector ([B^0, B^1, ... B^L]),
# # and then all equivariant basis.
# # What can be adjusted in its input are: (1) total polynomial degree; (2) correlation order; (3) largest L
# # (4) weight of the order of spherical harmonics; (5) specified radial basis
# function equivariant_luxchain_constructor(totdeg, ν, L; wL = 1, Rn = legendre_basis(totdeg))

# filter = RPE_filter_long(L)
# cgen = Rot3DCoeffs_long(L)

# Ylm = CYlmBasis(totdeg)

# spec1p = make_nlms_spec(simple_radial_basis(Rn), Ylm; totaldegree = totdeg, admissible = (br, by) -> br.n + wL * by.l <= totdeg)
# spec1p = sort(spec1p, by = (x -> x.n + x.l * wL))
# spec1pidx = getspec1idx(spec1p, simple_radial_basis(Rn).Radialspec, Ylm)

# # define sparse for n-correlations
# tup2b = vv -> [ spec1p[v] for v in vv[vv .> 0] ]
# default_admissible = bb -> length(bb) == 0 || sum(b.n for b in bb) + wL * sum(b.l for b in bb) <= totdeg

# specAA = gensparse(; NU = ν, tup2b = tup2b, filter = filter,
# admissible = default_admissible,
# minvv = fill(0, ν),
# maxvv = fill(length(spec1p), ν),
# ordered = true)

# spec = [ vv[vv .> 0] for vv in specAA if !(isempty(vv[vv .> 0]))]
# # map back to nlm
# spec_nlm = getspecnlm(spec1p, spec)

# C = _rpi_A2B_matrix(cgen, spec_nlm)

# # acemodel with lux layers
# bA = P4ML.PooledSparseProduct(spec1pidx)
# bAA = P4ML.SparseSymmProd(spec)

# # wrapping into lux layers
# l_Rn = P4ML.lux(Rn)
# l_Ylm = P4ML.lux(Ylm)
# l_bA = P4ML.lux(bA)
# l_bAA = P4ML.lux(bAA)

# # formming model with Lux Chain
# _norm(x) = norm.(x)

# l_xnx = Lux.Parallel(nothing; normx = WrappedFunction(_norm), x = WrappedFunction(identity))
# l_embed = Lux.Parallel(nothing; Rn = l_Rn, Ylm = l_Ylm)

# l1l2set = [(l1,l2) for l1 = 0:L for l2 = 0:L-l1]
# cgmat_set = [cgmatrix(l1,l2) for (l1,l2) in l1l2set ]

# _block(x,cgmat,l1,l2) = reshape.(Ref(cgmat) .* [ x[j][1:size(cgmat)[2]] for j = 1:length(x) ], 2l1+1, 2l2+1)
# _dropzero(x) = x[findall(X -> X>1e-12, norm.(x))]

# l_seperate = Lux.Parallel(nothing, [WrappedFunction(x -> _dropzero(_block(x,cgmat_set[i],l1l2set[i][1],l1l2set[i][2]))) for i = 1:length(cgmat_set)]... )

# # C - A2Bmap
# luxchain = Chain(xnx = l_xnx, embed = l_embed, A = l_bA , AA = l_bAA, BB = WrappedFunction(x -> C * x), blocks = l_seperate)#, rAA = WrappedFunction(ComplexF64))
# ps, st = Lux.setup(MersenneTwister(1234), luxchain)

return luxchain, ps, st
end

# This constructor builds a lux chain that maps a configuration to B^0, B^1, ... B^L first,
# and then deduces all equivariant basis.
# What can be adjusted in its input are: (1) total polynomial degree; (2) correlation order; (3) largest L
# (4) weight of the order of spherical harmonics; (5) specified radial basis
function equivariant_luxchain_constructor_new(totdeg, ν, L; wL = 1, Rn = legendre_basis(totdeg))
Ylm = CYlmBasis(totdeg)

spec1p = make_nlms_spec(simple_radial_basis(Rn), Ylm; totaldegree = totdeg, admissible = (br, by) -> br.n + wL * by.l <= totdeg)
spec1p = sort(spec1p, by = (x -> x.n + x.l * wL))
spec1pidx = getspec1idx(spec1p, simple_radial_basis(Rn).Radialspec, Ylm)

# define sparse for n-correlations
tup2b = vv -> [ spec1p[v] for v in vv[vv .> 0] ]
default_admissible = bb -> length(bb) == 0 || sum(b.n for b in bb) + wL * sum(b.l for b in bb) <= totdeg

# initialize C and spec_nlm
C = Vector{Any}(undef,L+1)
# spec_nlm = Vector{Any}(undef,L+1)
spec = Vector{Any}(undef,L+1)

# Spec = []

for l = 0:L
filter = RPE_filter(l)
cgen = Rot3DCoeffs(l)
specAA = gensparse(; NU = ν, tup2b = tup2b, filter = filter,
admissible = default_admissible,
minvv = fill(0, ν),
maxvv = fill(length(spec1p), ν),
ordered = true)

spec[l+1] = [ vv[vv .> 0] for vv in specAA if !(isempty(vv[vv .> 0]))]
# Spec = Spec ∪ spec

spec_nlm = getspecnlm(spec1p, spec[l+1])
C[l+1] = _rpi_A2B_matrix(cgen, spec_nlm)
end
# return C, spec, spec_nlm
# acemodel with lux layers
bA = P4ML.PooledSparseProduct(spec1pidx)

#
filter = RPE_filter_long(L)
specAA = gensparse(; NU = ν, tup2b = tup2b, filter = filter,
admissible = default_admissible,
minvv = fill(0, ν),
maxvv = fill(length(spec1p), ν),
ordered = true)

Spec = [ vv[vv .> 0] for vv in specAA if !(isempty(vv[vv .> 0]))]
dict = Dict([Spec[i] => i for i = 1:length(Spec)])
pos = [ [dict[spec[k][j]] for j = 1:length(spec[k])] for k = 1:L+1 ]
# @show typeof(pos)
# return C, spec, Spec

bAA = P4ML.SparseSymmProd(Spec)

# wrapping into lux layers
l_Rn = P4ML.lux(Rn)
l_Ylm = P4ML.lux(Ylm)
l_bA = P4ML.lux(bA)
l_bAA = P4ML.lux(bAA)

# formming model with Lux Chain
_norm(x) = norm.(x)

l_xnx = Lux.Parallel(nothing; normx = WrappedFunction(_norm), x = WrappedFunction(identity))
l_embed = Lux.Parallel(nothing; Rn = l_Rn, Ylm = l_Ylm)
l_seperate = Lux.Parallel(nothing, [WrappedFunction(x -> C[i] * x[pos[i]]) for i = 1:L+1]... )

function _vectorize(cc)
ccc = []
for i = 1:length(cc)
ccc = [ccc; cc[i]...]
end
return sparse(complex.(ccc))
end

function _block(x,l1,l2)
@assert l1+l2 <= length(x) - 1
init = iseven(l1+l2) ? 0 : 1
fea_set = init:2:l1+l2
A = Vector{Any}(undef,l1+l2+1)
for i = 0:l1+l2
if i in fea_set
A[i+1] = x[i+1]
else
A[i+1] = [zeros(ComplexF64,2i+1)]
end
end
cc = Iterators.product(A...) |> collect
c = [ _vectorize(cc[i]) for i = 1:length(cc) ]
return reshape.( Ref(cgmatrix(l1,l2)) .* c, 2l1+1, 2l2+1 )
end

# l_condensed = WrappedFunction(x -> _block(x,0,1))
l_condensed = Lux.Parallel(nothing, [WrappedFunction(x -> _block(x,l1,l2)) for l1 in 0:L for l2 in 0:L-l1]... )


# C - A2Bmap
luxchain = Chain(xnx = l_xnx, embed = l_embed, A = l_bA , AA = l_bAA, BB = l_seperate, inter = WrappedFunction(x -> [x[i] for i = 1:length(x)]), B = l_condensed)#, rAA = WrappedFunction(ComplexF64))
ps, st = Lux.setup(MersenneTwister(1234), luxchain)

return luxchain, ps, st
end
# return luxchain, ps, st
# end

# # This constructor builds a lux chain that maps a configuration to B^0, B^1, ... B^L first,
# # and then deduces all equivariant basis.
# # What can be adjusted in its input are: (1) total polynomial degree; (2) correlation order; (3) largest L
# # (4) weight of the order of spherical harmonics; (5) specified radial basis
# function equivariant_luxchain_constructor_new(totdeg, ν, L; wL = 1, Rn = legendre_basis(totdeg))
# Ylm = CYlmBasis(totdeg)

# spec1p = make_nlms_spec(simple_radial_basis(Rn), Ylm; totaldegree = totdeg, admissible = (br, by) -> br.n + wL * by.l <= totdeg)
# spec1p = sort(spec1p, by = (x -> x.n + x.l * wL))
# spec1pidx = getspec1idx(spec1p, simple_radial_basis(Rn).Radialspec, Ylm)

# # define sparse for n-correlations
# tup2b = vv -> [ spec1p[v] for v in vv[vv .> 0] ]
# default_admissible = bb -> length(bb) == 0 || sum(b.n for b in bb) + wL * sum(b.l for b in bb) <= totdeg

# # initialize C and spec_nlm
# C = Vector{Any}(undef,L+1)
# # spec_nlm = Vector{Any}(undef,L+1)
# spec = Vector{Any}(undef,L+1)

# # Spec = []

# for l = 0:L
# filter = RPE_filter(l)
# cgen = Rot3DCoeffs(l)
# specAA = gensparse(; NU = ν, tup2b = tup2b, filter = filter,
# admissible = default_admissible,
# minvv = fill(0, ν),
# maxvv = fill(length(spec1p), ν),
# ordered = true)

# spec[l+1] = [ vv[vv .> 0] for vv in specAA if !(isempty(vv[vv .> 0]))]
# # Spec = Spec ∪ spec

# spec_nlm = getspecnlm(spec1p, spec[l+1])
# C[l+1] = _rpi_A2B_matrix(cgen, spec_nlm)
# end
# # return C, spec, spec_nlm
# # acemodel with lux layers
# bA = P4ML.PooledSparseProduct(spec1pidx)

# #
# filter = RPE_filter_long(L)
# specAA = gensparse(; NU = ν, tup2b = tup2b, filter = filter,
# admissible = default_admissible,
# minvv = fill(0, ν),
# maxvv = fill(length(spec1p), ν),
# ordered = true)

# Spec = [ vv[vv .> 0] for vv in specAA if !(isempty(vv[vv .> 0]))]
# dict = Dict([Spec[i] => i for i = 1:length(Spec)])
# pos = [ [dict[spec[k][j]] for j = 1:length(spec[k])] for k = 1:L+1 ]
# # @show typeof(pos)
# # return C, spec, Spec

# bAA = P4ML.SparseSymmProd(Spec)

# # wrapping into lux layers
# l_Rn = P4ML.lux(Rn)
# l_Ylm = P4ML.lux(Ylm)
# l_bA = P4ML.lux(bA)
# l_bAA = P4ML.lux(bAA)

# # formming model with Lux Chain
# _norm(x) = norm.(x)

# l_xnx = Lux.Parallel(nothing; normx = WrappedFunction(_norm), x = WrappedFunction(identity))
# l_embed = Lux.Parallel(nothing; Rn = l_Rn, Ylm = l_Ylm)
# l_seperate = Lux.Parallel(nothing, [WrappedFunction(x -> C[i] * x[pos[i]]) for i = 1:L+1]... )

# function _vectorize(cc)
# ccc = []
# for i = 1:length(cc)
# ccc = [ccc; cc[i]...]
# end
# return sparse(complex.(ccc))
# end

# function _block(x,l1,l2)
# @assert l1+l2 <= length(x) - 1
# init = iseven(l1+l2) ? 0 : 1
# fea_set = init:2:l1+l2
# A = Vector{Any}(undef,l1+l2+1)
# for i = 0:l1+l2
# if i in fea_set
# A[i+1] = x[i+1]
# else
# A[i+1] = [zeros(ComplexF64,2i+1)]
# end
# end
# cc = Iterators.product(A...) |> collect
# c = [ _vectorize(cc[i]) for i = 1:length(cc) ]
# return reshape.( Ref(cgmatrix(l1,l2)) .* c, 2l1+1, 2l2+1 )
# end

# # l_condensed = WrappedFunction(x -> _block(x,0,1))
# l_condensed = Lux.Parallel(nothing, [WrappedFunction(x -> _block(x,l1,l2)) for l1 in 0:L for l2 in 0:L-l1]... )


# # C - A2Bmap
# luxchain = Chain(xnx = l_xnx, embed = l_embed, A = l_bA , AA = l_bAA, BB = l_seperate, inter = WrappedFunction(x -> [x[i] for i = 1:length(x)]), B = l_condensed)#, rAA = WrappedFunction(ComplexF64))
# ps, st = Lux.setup(MersenneTwister(1234), luxchain)

# return luxchain, ps, st
# end
Loading

0 comments on commit bf3a523

Please sign in to comment.