张量网络碎碎念：CGC

在本系列 上一篇文章 中，我介绍了张量网络的一些基础概念。其中很大一部分来自 github 上一个教程。事实上，该教程的大部分内容来自 e3nn 官网。
除了上篇文章介绍的一些可视化技巧，官网还提供了其他一些可视化模块。使用这些功能能使我们更深入理解张量网络中的种种细节。其中，本人觉得最实用的功能就是张量积通路的可视化。（事实上，我只看懂这一个hhh）
下面就张量积为例介绍：
e3nn 中的张量积形式是 CGC 矩阵实现的，用户无需关注计算细节，只需提供输入向量及其表征即可：

import os

import matplotlib.pyplot as plt
from e3nn import o3


irreps1 = o3.Irreps("1x1e")
irreps2 = o3.Irreps("1x1e")

tp = o3.FullTensorProduct(irreps1, irreps2)
print(tp)
tp.visualize()
plt.show()

input1 = irreps1.randn(-1)
input2 = irreps2.randn(-1)
results = tp(input1, input2)
print(results)

其中

tp.visualize()

是 e3nn 内置的可视化通道，本例中，我们没有对输出张量的格式做限定，所以输出张量最低价是 1-1=0，最高价是 1+1=2

我们随机初始化符合输入表征的两个输入向量，并将其填入预先划定好的 tp 通道，即可得到最终的张量积结果：

input1 = irreps1.randn(-1)
input2 = irreps2.randn(-1)
results = tp(input1, input2)
print(results)

可以看到，使用 e3nn ，我们轻松完成了一次张量积。事实上，在张量积背后，e3nn 是使用 CGC 完成的计算。这导致张量积一直以来造成性能瓶颈，据说是 O(Lmax^6)，导致诸多张量网络不敢增加 Lmax 值。
如何优化张量积计算，解决性能瓶颈，目前主流的算法是绕开 CGC 模式，从 SO3 降到 SO2 (eSCN)。
本文将先关注 CGC，试图回答，CGC 做了什么？
首先，CGC 公式如下：

等式左边：
x, y 分别是输入张量，做直积。但是直积的结果有很多项，有很多的角动量，每个角动量下面有磁动量。
具体来说，

在确定了角动量后，磁动量的取值范围是

在本例中，两输入均为 1e，所以张量积结果的角动量在 0，1，2 间，如上图可视化所示。
每一个角动量，对应不同的磁动量，所以等式左边一共会有 1+3+5=9 种 (l3, m3) 组合。
在确定了 (l3, m3) 以后（9选1），由于

等式左边固定，等式右边仅剩下 m1, m2 两个可变参数（取值范围跟 l1, l2 有关，本例中是 0, +1, -1），这也是大大的求和符号的存在意义：
对于特定的 (l3, m3) ，我们需要遍历 m1, m2，在遍历的过程中，我们根据 m1, m2 索引 x,y，将其内积，最后再根据 m1, m2 以及前面固定的 (l3, m3, l1, l2) ，我们能求出 CGC 系数，将其和内积结果拼到一起即可。
最终，我们将所有内积结果加在一起，即可得到等式左边的值，的其中一项。我们重复固定 (l3, m3) ，最终能求出全部张量结果。
为了验证这一想法，我们使用 sympy 里的 cgc 模块，求出 cgc ，再进行拼凑。
全部代码如下：

import os

import matplotlib.pyplot as plt
from e3nn import o3
from sympy import S
from sympy.physics.quantum.cg import CG


def clebsch_gordan_coefficients(l1, m1, l2, m2, l, m):
    cg = CG(S(l1), S(m1), S(l2), S(m2), S(l), S(m)).doit()
    return float(cg)


irreps1 = o3.Irreps("1x1e")
irreps2 = o3.Irreps("1x1e")

tp = o3.FullTensorProduct(irreps1, irreps2)
print(tp)
tp.visualize()
plt.show()

input1 = irreps1.randn(1, -1)
input2 = irreps2.randn(1, -1)
results = tp(input1, input2)

# Possible values of total angular momentum J and its projection M
total_angular_momenta = [(0, 0), (1, 1), (1, 0), (1, -1), (2, 0), (2, 2), (2, 1), (2, -1), (2, -2)]
all_cgc_results = []
for idx, (l3, m3) in enumerate(total_angular_momenta):
    sub_results = []
    for idx1, m1 in enumerate([-1, 1, 0]):
        for idx2, m2 in enumerate([-1, 1, 0]):
            # print(f"x(l1=1, m1={m1}), y(l1=1, m2={m2}) for l3={l3}, m3={m3}")
            coefficient = clebsch_gordan_coefficients(l1=1, l2=1, m1=m1, m2=m2, l=l3, m=m3)
            # print(f"Clebsch-Gordan Coefficient for (l1=1, m1={m1}, l2=1, m2={m2}, l={l3}, m={m3}): {coefficient}")
            a_sub_result = coefficient * input1[0][m1] * input2[0][m2]
            sub_results.append(a_sub_result)
    sum_sub_results = sum(sub_results)
    print(f'The l={l3}, m={m3} has been iterated. \n'
          f'The tp results provided by e3nn is {results[0][idx]}.\n'
          f'The tp results computed with assistance of CGC is {sum_sub_results}')
    all_cgc_results.append(round(float(sum_sub_results), 4))
print('\nThe tp results:\n')
print(results)
print('\nThe cgc results:\n')
print(all_cgc_results)

结果如下：

The tp results:

tensor([[ 0.7047,  1.3212, -0.5690,  0.1757,  0.5107,  0.6395,  1.0796,  1.0922,
         -0.2554]])

The cgc results:

[0.8071, -0.1757, 1.3212, 0.569, 0.7504, 1.2884, 0.6395, 0.5107, -0.2145]

其中 9 个数中 5 个数字的绝对值是完全吻合的，顺序是颠倒的，未能全部复现成功可能因为，磁动量的顺序在不同的程序包中不一致，以及其他索引，倍数差异等问题。尽管如此，我们复现出了绝对值完全一致的结果，证明了 e3nn 内部的黑箱是按照 CGC 模式进行计算的。