WIP : Naive 1-Corr Purification

src/Polynomials4ML.jl

@@ -84,6 +84,9 @@ include("ace/symmprod_dag_kernels.jl")
 include("ace/simpleprodbasis.jl")
 include("ace/sparsesymmprod.jl")
 
+# experimental 
+include("ace/pure2b.jl")
+
 # some nice utility functions to generate basis sets and other things  
 include("utils/utils.jl")
 

src/ace/pure2b.jl

@@ -0,0 +1,89 @@
+
+struct Pure2B{TA, TAA} <: AbstractP4MLTensor
+   abasis::TA 
+   aabasis::TAA 
+   # --------------
+   @reqfields
+end
+
+function Pure2B(abasis, aabasis)
+   return Pure2B(abasis, aabasis, _make_reqfields()...)
+end
+
+Base.length(basis::Pure2B) = length(basis.aabasis)
+
+Base.show(io::IO, basis::Pure2B{TA, TAA}) where {TA, TAA} = 
+      print(io, "Pure2B($(TA.name.name), $(TAA.name.name))")
+
+# -------------- evaluation interfaces
+
+_valtype(basis::Pure2B, args...) = _valtype(basis.abasis, args...)
+
+_gradtype(basis::Pure2B, args...) = _gradtype(basis.abasis, args...)
+
+_generate_input(basis::Pure2B; nX = rand(5:15)) = 
+   _generate_input(basis.abasis; nX = nX)
+
+function _generate_batch(basis::Pure2B, args...; kwargs...) 
+   error("Pure2B is not implemented for batch inputs")
+end
+
+function whatalloc(::typeof(evaluate!), basis::Pure2B, 
+                   BB::Tuple{Matrix{T1}, Matrix{T2}}) where {T1, T2}
+   VT = promote_type(T1, T2) 
+   return (VT, length(basis))
+end
+
+# ----------------------- evaluation kernels 
+
+# specialized implementation for two inputs only
+# the prototype is Rnl, Ylm. The comments in this function are copied 
+# from the AA basis implementation. 
+
+function evaluate!(AA, basis::Pure2B, BB::Tuple{TB1, TB2}) where {TB1, TB2}
+   aspec = basis.abasis.spec
+   nodes = basis.aabasis.nodes
+   has0 = basis.aabasis.has0
+   num1 = basis.aabasis.num1
+   @assert length(AA) >= basis.aabasis.numstore
+
+   B1, B2 = BB
+   @assert size(B1, 1) == size(B2, 1) 
+   nX = size(B1, 1)
+
+   TAA = eltype(AA)
+
+   if has0
+      error("Pure2B does not support has0 == true")
+   end 
+
+   PHI = zeros(TAA, nX, length(nodes))
+
+   # Stage-1: copy the 1-particle basis into AA
+   # note this entirely ignores the spec / nodes. It is implicit in the 
+   # definitions and orderings
+   @inbounds for i = 1:num1
+      # AA[has0+i] = A[i]
+      n1, n2 = aspec[i] 
+      ai = zero(TAA) 
+      for j = 1:nX 
+         PHI[j, i] = ϕ_ij = B1[j, n1] * B2[j, n2]
+         ai += ϕ_ij
+      end
+      AA[i] = ai 
+   end
+
+   # Stage-2: go through the dag and store the intermediate results we need
+   @inbounds for i = (num1+has0+1):length(nodes)
+      n1, n2 = nodes[i]
+      # AA[i] = AA[n1] * AA[n2]
+      aai = zero(TAA)
+      for j = 1:nX 
+         PHI[j, i] = ϕ_ij = PHI[j, n1] * PHI[j, n2]
+         aai += ϕ_ij
+      end 
+      AA[i] = aai
+   end
+
+   return AA
+end

test/ace/test_pure2b.jl

@@ -0,0 +1,84 @@
+using Test, Polynomials4ML, ChainRulesCore
+using Polynomials4ML: SparseSymmProdDAG, 
+                      evaluate, evaluate_d, evaluate_dd
+using Polynomials4ML.Testing: println_slim, print_tf, generate_SO2_spec
+using Random
+
+using ACEbase.Testing: fdtest, dirfdtest
+
+P4ML = Polynomials4ML
+
+maxn = 8; maxl = 4; 
+aspec = [ (n, l) for n = 1:maxn for l = 1:maxl ]
+abasis = PooledSparseProduct(aspec)
+lena = length(aspec)
+aaspec = [ [ [ k, ] for k = 1:lena ]; 
+           [ [k1, k2]  for k1 = 1:lena for k2 = k1:lena ]; 
+           [ [k1, k2, k3] for k1 = 1:lena for k2 = k1:lena for k3 = k2:lena ] ]
+aaspec = filter( bb -> sum(bb) <= 20, aaspec)
+aabasis = SparseSymmProdDAG(aaspec)
+basis = P4ML.Pure2B(abasis, aabasis)
+
+I1 = findall(length.(aaspec) .== 1)
+I2 = findall(length.(aaspec) .== 2)
+I3 = findall(length.(aaspec) .== 3)
+
+##
+
+nX = 12
+R = randn(nX, maxn) 
+Y = randn(nX, maxl) 
+A = basis.abasis((R, Y))
+AA = basis.aabasis(A)
+ϕ = evaluate(basis, (R, Y))
+AA2 = AA - ϕ
+
+##
+
+@info("confirm permutation invariance")
+p = shuffle(1:nX)
+Rp = R[p, :]
+Yp = Y[p, :]
+Ap = basis.abasis((Rp, Yp))
+AAp = basis.aabasis(Ap)
+ϕp = evaluate(basis, (Rp, Yp))
+AA2p = AAp - ϕp
+println_slim(@test AA2 ≈ AA2p) 
+
+##
+
+# to show that there is no 2B we can write it explicitly as sum over clusters 
+# the actual 2B (1-corr) terms should just be zero now 
+@info("Confirm purified 1-corr")
+println_slim(@test all(abs.(AA2[I1]) .<= 1e-14))
+
+@info("Confirm purified 2-corr")
+# for 3B / 2-Corr we sum over all clusters 
+AA_cl2 = zeros(size(AA2[I2]))
+K1 = [ bb[1] for bb in aaspec[I2] ]
+K2 = [ bb[2] for bb in aaspec[I2] ]
+for j1 = 1:nX, j2 = 1:nX 
+   if j1 != j2 
+      ϕ1 = [ R[j1, n] * Y[j1, l] for (n, l) in aspec ]
+      ϕ2 = [ R[j2, n] * Y[j2, l] for (n, l) in aspec ]
+      AA_cl2 += ϕ1[K1] .* ϕ2[K2]
+   end
+end
+println_slim(@test AA_cl2 ≈ AA2[I2])
+
+@info("Confirm purified 3-corr")
+# same for 4B / 3-corr
+AA_cl3 = zeros(size(AA2[I3]))
+K1 = [ bb[1] for bb in aaspec[I3] ]
+K2 = [ bb[2] for bb in aaspec[I3] ]
+K3 = [ bb[3] for bb in aaspec[I3] ]
+for j1 = 1:nX, j2 = 1:nX, j3 = 1:nX
+   if !(j1 == j2 == j3)
+      ϕ1 = [ R[j1, n] * Y[j1, l] for (n, l) in aspec ]
+      ϕ2 = [ R[j2, n] * Y[j2, l] for (n, l) in aspec ]
+      ϕ3 = [ R[j3, n] * Y[j3, l] for (n, l) in aspec ]
+      AA_cl3 += ϕ1[K1] .* ϕ2[K2] .* ϕ3[K3]
+   end
+end
+println_slim(@test AA_cl3 ≈ AA2[I3])
+