From 64d8b4495f46eb3163909c7c6af29d19fdaa9dd4 Mon Sep 17 00:00:00 2001
From: dussong <genevieve.dusson@univ-fcomte.fr>
Date: Thu, 31 Oct 2024 10:55:12 +0100
Subject: [PATCH] test RPI - debug

---
 Project.toml           |   1 +
 src/O3_alternative.jl  | 176 ++++++++++++++++++++++++++++++++++++++++-
 test/test_RI_basis.jl  |  66 +++++++++++++++-
 test/test_RPI_basis.jl | 126 +++++++++++------------------
 4 files changed, 282 insertions(+), 87 deletions(-)

diff --git a/Project.toml b/Project.toml
index 6eef21a..b820cb9 100644
--- a/Project.toml
+++ b/Project.toml
@@ -11,6 +11,7 @@ Polynomials4ML = "03c4bcba-a943-47e9-bfa1-b1661fc2974f"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpheriCart = "5caf2b29-02d9-47a3-9434-5931c85ba645"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 
 [compat]
 StaticArrays = "1.5"
diff --git a/src/O3_alternative.jl b/src/O3_alternative.jl
index 37d10b0..5ade884 100644
--- a/src/O3_alternative.jl
+++ b/src/O3_alternative.jl
@@ -4,7 +4,7 @@ using PartialWaveFunctions
 using Combinatorics
 using LinearAlgebra
 
-export re_basis_new
+export re_basis_new, ri_basis_new, ind_corr_s1, ind_corr_s2, MatFmi, ML0
 
 function CG(l,m,L,N) 
     M=m[1]+m[2]
@@ -50,6 +50,30 @@ function SetLl0(l,N)
     return set
 end
 
+function SetLl(l,N,L)
+    set = Vector{Int64}[]
+    for k in abs(l[1]-l[2]):l[1]+l[2]
+        push!(set, [0; k])
+    end
+    for k in 3:N-1
+        setL=set
+        set=Vector{Int64}[]
+        for a in setL
+            for b in abs(a[k-1]-l[k]):a[k-1]+l[k]
+                push!(set, [a; b])
+            end
+        end
+    end
+    setL=set
+    set=Vector{Int64}[]
+    for a in setL
+        if (abs.(a[N-1]-l[N]) <= L)&&(L <= (a[N-1]+l[N]))
+            push!(set, [a; L])
+        end
+    end
+    return set
+end
+
 # Function that computes the set ML0
 function ML0(l,N)
     setML = [[i] for i in -abs(l[1]):abs(l[1])]
@@ -70,7 +94,29 @@ function ML0(l,N)
     return setML0
 end
 
-function re_basis_new(l)
+# Function that computes the set ML (relative to equivariance L)
+function ML(l,N,L)
+    setML = [[i] for i in -abs(l[1]):abs(l[1])]
+    for k in 2:N-1
+        set = setML
+        setML = Vector{Int64}[]
+        for m in set
+            append!(setML, [m; lk] for lk in -abs(l[k]):abs(l[k]) )
+        end
+    end
+    setML0=Vector{Int64}[]
+    for m in setML
+        s=sum(m)
+        for mn in -L-s:L-s
+            if  abs(mn) < abs(l[N])+1
+                push!(setML0, [m; mn])
+            end
+        end
+    end
+    return setML0
+end
+
+function ri_basis_new(l)
     N=size(l,1)
     L=SetLl0(l,N)
     r=size(L,1)
@@ -89,4 +135,130 @@ function re_basis_new(l)
         end
     end
     return U,M
+end
+
+function re_basis_new(l,L)
+    N=size(l,1)
+    Ll=SetLl(l,N,L)
+    r=size(Ll,1)
+    if r==0 
+        return zeros(Float64, 0, 0)
+    else 
+        setML0=ML(l,N,L)
+        sizeML0=length(setML0)
+        U=zeros(Float64, r, sizeML0)
+        M = Vector{Int64}[]
+        for (j,m) in enumerate(setML0)
+            push!(M,m)
+            for i in 1:r
+                U[i,j]=CG(l,m,Ll[i],N)
+            end
+        end
+    end
+    return U,M
+end
+
+
+# Function that computes the permutations that let n and l invariant
+function Snl(N,n,l)
+    if n==n[1]*ones(N)
+        if l==l[1]*ones(N)
+            return permutations(1:N)
+        end
+    end
+    if N==1
+        return Set([[1]])
+    elseif (n[N-1],l[N-1])!=(n[N],l[N])
+        S=Set()
+        Sn=Snl(N-1,n[1:N-1],l[1:N-1])
+        for x in Sn
+            append!(x,[N])
+            union!(S,Set([x]))
+        end
+    else
+        S=Set()
+        k=N
+        while (n[k-1],l[k-1])==(n[k],l[k]) && k>2
+            k-=1
+        end
+        if k==2 && (n[1],l[1])==(n[2],l[2])
+            return Set(permutations(1:N))
+        else
+            Sn=Snl(k-1,n[1:k-1],l[1:k-1])
+            for x in Sn
+                for s in Set(permutations(k:N))
+                    y=copy(x)
+                    append!(y,s)
+                    union!(S,Set([y]))
+                end
+            end
+        end
+    end
+    return S
+end
+
+
+#Function that computes the set of classes using the set Ml0 and the possible permutations
+function class(setML0,sigma,N,l)
+    setclass=Vector{Vector{Int64}}[]
+    pop!(setML0,zeros(Int64,N))
+    while setML0!=Set()
+        x=pop!(setML0)
+        p=[x]
+        for s in sigma
+            y=x[s]
+            if y in setML0
+                append!(p,[y])
+                pop!(setML0,y)
+            end
+        end
+        append!(setclass,[p])
+    end
+    setclasses=Vector{Vector{Int64}}[]
+    for x in setclass
+        for y in setclass
+            if x==y
+                if minimum(x)==minimum(-x)
+                    if iseven(sum(l))
+                        append!(setclasses,[x])
+                    end
+                end
+            elseif minimum(x)==minimum(-y)
+                if y<x
+                    append!(setclasses,[x])
+                end
+            end
+        end
+    end
+    if iseven(sum(l))
+        append!(setclasses,[[zeros(N)]])
+    end
+    setclasses
+end
+
+
+
+# Function that computes the matrix ( f(m,i) )
+function MatFmi(n,l)
+    N=size(l,1)
+    L=SetLl0(l,N)
+    r=size(L,1)
+    if r==0 
+        return zeros(Float64, 0, 0)
+    else 
+        ML00 = ML0(l,N)
+        setML0=Set(ML00)
+        sigma = Snl(N,n,l)
+        setclass=class(setML0,sigma,N,l)
+        sizeML0=length(setclass)
+        Matrix=zeros(Float64, r, sizeML0)
+        for i in 1:r
+            for j in 1:sizeML0
+                for m in setclass[j]
+                    Matrix[i,j]+=CG(l,m,L[i],N)
+                end
+            end
+        end
+    end
+    return Matrix, ML00
 end
\ No newline at end of file
diff --git a/test/test_RI_basis.jl b/test/test_RI_basis.jl
index 1f13b77..b1a5599 100644
--- a/test/test_RI_basis.jl
+++ b/test/test_RI_basis.jl
@@ -57,14 +57,14 @@ end
 # for the moment the code with generalized CG only works with L=0
 L = 0
 cc = Rot3DCoeffs(L)
-ll = SA[2,2,2,3,3]
+ll = SA[1,1,2,2,4]
 
 # version with svd
 @time coeffs1, MM1 = O3.re_basis(cc, ll)
 nbas = size(coeffs1, 1)
 
 #version with gen CG coefficients
-@time coeffs2, MM2 = re_basis_new(ll)
+@time coeffs2, MM2 = ri_basis_new(ll)
 
 # simple test on size
 @test size(coeffs1) == size(coeffs2)
@@ -85,7 +85,7 @@ coeffsp2 = coeffs2[:,P2]
 @test rank(coeffsp2) == size(coeffsp2,1)
 
 # check that the coef span the same space - test fails
-@test nbas == rank([coeffsp; coeffsp2], rtol = 1e-12)
+@test nbas == rank([coeffsp1; coeffsp2], rtol = 1e-12)
 
 
 Rs = [rand_sphere() for _ in 1:length(ll)]
@@ -122,4 +122,62 @@ end
 
 # check that functions span same space
 rk = rank([A1;A2]; rtol = 1e-12)  
-@test rk == nbas
\ No newline at end of file
+@test rk == nbas
+
+# ----------------------------
+# Extension to equivariance 
+# ----------------------------
+
+# Test SetLl0
+N = 5
+l = rand(1:10, N)
+SL1 = RepLieGroups.SetLl(l,N,0)
+SL2 = RepLieGroups.SetLl0(l,N)
+@test SL1 == SL2
+
+# Test SetML 
+N = 5
+l = rand(1:10, N)
+SM1 = RepLieGroups.ML(l,N,0)
+SM2 = RepLieGroups.ML0(l,N)
+@test SM1 == SM2
+
+ 
+# # Version CO
+# L = 1
+# cc = Rot3DCoeffs(L)
+# ll = SA[1,1,2,3,4]
+
+# # version with svd
+# @time coeffs1, MM1 = O3.re_basis(cc, ll)
+# nbas = size(coeffs1, 1)
+
+# # Version GD
+# @time coeffs2, MM2 = re_basis_new(ll,L)
+
+# # simple test on size
+# @test size(coeffs1) == size(coeffs2)
+# @test size(MM1) == size(MM2)
+
+# P1 = sortperm(MM1)
+# P2 = sortperm(MM2)
+# MMsorted1 = MM1[P1]
+# MMsorted2 = MM2[P2]
+# # check that same mm values
+# @test MMsorted1 == MMsorted2
+
+# coeffsp1 = coeffs1[:,P1]
+# coeffsp2 = coeffs2[:,P2]
+
+# # test that full rank
+# @test rank(coeffsp1) == size(coeffsp1,1)
+# @test rank(coeffsp2) == size(coeffsp2,1)
+
+# # check that the coef span the same space 
+# @test nbas == rank([coeffsp1; coeffsp2], rtol = 1e-12)
+
+# Rs = [rand_sphere() for _ in 1:length(ll)]
+# Q = rand_rot() 
+# QRs = [Q*Rs[i] for i in 1:length(Rs)]
+# fRs1 = [ f(Rs, q; coeffs=coeffs1, MM=MM1, ll=ll) for q = 1:nbas ]
+# fRs1Q = [ f(QRs, q; coeffs=coeffs1, MM=MM1, ll=ll) for q = 1:nbas ]
diff --git a/test/test_RPI_basis.jl b/test/test_RPI_basis.jl
index 6ac48fc..ec143cd 100644
--- a/test/test_RPI_basis.jl
+++ b/test/test_RPI_basis.jl
@@ -1,9 +1,10 @@
 
 using SpheriCart, StaticArrays, LinearAlgebra, RepLieGroups, WignerD, 
       Combinatorics
-using Rotations
-using RepLieGroups.O3: Rot3DCoeffs
+using RepLieGroups.O3: Rot3DCoeffs, Rot3DCoeffs_real
 O3 = RepLieGroups.O3
+using Test
+
 
 function eval_cY(rbasis::SphericalHarmonics{LMAX}, 𝐫) where {LMAX}  
    Yr = rbasis(𝐫)
@@ -30,112 +31,75 @@ function rand_sphere()
    return u / norm(u) 
 end
 
+function rand_radial(rmax) 
+   return rmax*rand()
+end
+
 function rand_rot() 
    K = @SMatrix randn(3,3)
    return exp(K - K') 
 end
 
-
-##
-# CASE 2: 4-correlations, L = 0 (revisited)
-L = 0
-cc = Rot3DCoeffs(0)
-# now we fix an ll = (l1, l2, l3) triple ask for all possible linear combinations 
-# of the tensor product basis   Y[l1, m1] * Y[l2, m2] * Y[l3, m3] * Y[l4, m4]
-# that are invariant under O(3) rotations.
-ll = SA[3, 3, 2, 2, 2]
-coeffs, MM = O3.re_basis(cc, ll)
-nbas = size(coeffs, 1)
-# coeffs = nbasis x length(MM) matrix 
-#     MM = vector of (m1, m2, m3, m4) tuples 
-# for 4-correlations and higher, the number of possible couplings can be 
-# greater than one. An interesting result of this is that even if those 
-# basis functions as tensor products of Ylms are linearly independent, they 
-# need no longer be linearly independent once we impose permutation-invariance. 
-
 """
 This implements a permutation-invariant and O(3) invariant function of 
 4 (or more) variables. 
 """
-function f(Rs, q::Integer; coeffs=coeffs, MM=MM)
-   real_basis = SphericalHarmonics(3)
+function f(Rs, q::Integer; coeffs=coeffs, MM=MM, ll=ll)
+   N = length(coeffs[1])
+   Lmax = maximum(ll)
+   real_basis = SphericalHarmonics(Lmax)
    Y = [ eval_cY(real_basis, 𝐫) for 𝐫 in Rs ]
    A = sum(Y)   # this is a permutation-invariant embedding
    out = zero(eltype(A))
    for (c, mm) in zip(coeffs[q, :], MM)
-      ii = [SpheriCart.lm2idx(ll[α], mm[α]) for α in 1:4]
+      ii = [SpheriCart.lm2idx(ll[α], mm[α]) for α in 1:N]
       out += c * prod(A[ii])
    end
    return real(out)
 end
 
-
-Rs = [rand_sphere() for _ in 1:4]
-[ f(Rs, q) for q = 1:nbas ]
-
 # we can look for linear independence... 
-
-function f_rand_batch(coeffs, MM, ntest) 
+function f_rand_batch(coeffs, MM, ntest, ll) 
+   N = length(MM[1])
    nbas = size(coeffs, 1)
    A = zeros(nbas, ntest)
    for i = 1:ntest 
-      Rs = [rand_sphere() for _ in 1:4]
+      Rs = [rand_sphere() for _ in 1:N]
       for q = 1:nbas
-         A[q, i] = f(Rs, q; coeffs = coeffs)
+         A[q, i] = f(Rs, q; coeffs = coeffs, MM=MM, ll=ll)
       end
    end
    return A
 end
 
-A = f_rand_batch(coeffs, MM, 1_000)
-
-# the rank of this matrix is only 3, not nbas = 5!
-rk = rank(A)  
-# 3 
-
-# this is in fact very clear from the SVD 
-svdvals(A)
-# 5-element Vector{Float64}:
-#  25.568548451339442
-#   6.754493909270821
-#   5.45667827344765
-#   5.398138581058422e-15
-#   1.6731699590390224e-15
-
 ##
-# In ACE we use a semi-analytic construction to make the basis functions 
-# linearly indepdendent. 
-# in a full ACE code, this is a bit more complex since n channels are added 
-# to the story. This can be found here: 
-# https://github.com/ACEsuit/ACE1.jl/blob/8ac52d2128241a01d8b9a036f41b1d5106cbeb07/src/rpi/rotations3d.jl#L334
-# https://github.com/ACEsuit/ACE1.jl/blob/8ac52d2128241a01d8b9a036f41b1d5106cbeb07/src/rpi/rotations3d.jl#L348
-
-# For simplicity, we can just use the SVD to construct a 
-# linearly independent basis.
-
+# CASE 2: 4-correlations, L = 0 (revisited)
+L = 0
+cc = Rot3DCoeffs(L)
+# now we fix an ll = (l1, l2, l3) triple ask for all possible linear combinations 
+# of the tensor product basis   Y[l1, m1] * Y[l2, m2] * Y[l3, m3] * Y[l4, m4]
+# that are invariant under O(3) rotations.
+ll = SA[1,2,2,2,3]
+N = length(ll)
+nn = ones(Int64, N)# for the moment, nn has to be only ones
+@assert length(ll) == length(nn)
+@time coeffs1, MM1 = O3.re_basis(cc, ll)
+nbas_ri1 = size(coeffs1, 1)
+
+@time A = f_rand_batch(coeffs1,MM1, 1_000, ll) #this should depend on nn
+rk1 = rank(A, rtol = 1e-12) #does not correspond to what is in our paper
 U, S, V = svd(A)
-coeffs_ind = Diagonal(S[1:rk]) \ (U[:, 1:rk]' * coeffs)
-
-##
-# now let's re-compute A 
-
-ntest = 1_000
-A_ind = zeros(rk, ntest)
-for i = 1:ntest 
-   Rs = [rand_sphere() for _ in 1:4]
-   for q = 1:rk
-      A_ind[q, i] = f(Rs, q; coeffs = coeffs_ind)
-   end
-end
-
-# the rank of this matrix is 3, which is what we expect
-rk_ind = rank(A_ind)  
-# 3 
-
-# And we even scaled the coupling coeffs so that the singular values are 
-# now all close to 1. 
-svdvals(A_ind)
-# 3-element Vector{Float64}:
-#  0.997130965666211
-#  0.8853910366671397
-#  0.8668533817486788
\ No newline at end of file
+coeffs_ind1 = Diagonal(S[1:rk1]) \ (U[:, 1:rk1]' * coeffs1)
+
+# Version GD
+@time coeffs_rpi, MM_rpi = MatFmi(nn,ll)
+@show size(coeffs_rpi)
+@time coeffs2, MM2 = ri_basis_new(ll)
+
+rk2 = rank(coeffs_rpi,rtol = 1e-12)
+@test rk1 == rk2 #error here
+U, S, V = svd(coeffs_rpi)
+coeffs_ind2 = Diagonal(S[1:rk2]) \ (U[:, 1:rk2]' * coeffs2)
+# multiset_permutations([-1 -1 1 1], 4)
+# sort([[1, 1], [2, 3], [2, 2], [1, 4]])
+@test rank([coeffs_ind1;coeffs_ind2]) == rk1