使用 einsum 和多个数组的平均值进行加速

Question

给定两个矩阵 A 和 B，每个矩阵有 4 个索引，A[i,j,n_i,m], B[i,j,n_i,m]，我需要对前两个索引执行它们的张量积，而保留剩下的两个。对于此任务，我使用 einsum。此外，我需要对最后一个索引的结果进行平均，为此我只对特定轴使用 .mean() 。

C=np.einsum('ij...,kl...->ikjl...',A[:,:,:,:],B[:,:,:,:]).reshape(1024,1024,len(t_list),Mmax).mean(axis=3)   

虽然这可行，但我想知道是否有任何最佳方法可以做到这一点，因为现在这很慢。恐怕我的瓶颈是矩阵的大小：虽然原始的 A 和 B 矩阵的形状为 (32,32,1000,10000)，但生成的矩阵是 C=(1024,1024,1000)，这是相当昂贵的.

但是，使用 einsum 然后使用 .mean(axis=3) 求平均值仍然比循环 Mmax、创建数组然后取平均值要快得多。

我应该说我真正追求的是减少的矩阵，或者这个矩阵的部分痕迹，即

from qutip import *
dim=2*np.ones(2*5)
for n_i,n in enumerate(t_list):
    C_s=C[:,:,n_i]
    Qobj(C_s,dims=[dim,dim]).ptrace([0,1])

如果可能的话，我想做这个操作而不必直接分配C，我认为这是整个过程变慢的原因。

Answer 1

如评论中所述，您正在复制不必要的输入数组。此外，正如您所建议的，可以直接使用

np.einsum

计算外积的部分迹，而无需分配

C

。虽然它不会帮助代码可读性......

查看

ptrace

的实现和您传递给它的参数（也许下次尝试自己将这样的外部函数提炼成它的 Numpy 等价物 - 在某些上下文中可能更容易理解），看来您将结束向上调用

_ptrace_dense

（在这里），特别是该函数的其他部分：

rhomat = np.trace(Q.full()
    .reshape(rd + rd)
    .transpose(qtrace + [nd + q for q in qtrace] +
                sel + [nd + q for q in sel])
    .reshape([np.prod(dtrace, dtype=np.int32),
            np.prod(dtrace, dtype=np.int32),
            np.prod(dkeep, dtype=np.int32),
            np.prod(dkeep, dtype=np.int32)]))

如您所见，这会进行一些复杂的重塑/换位，然后获取结果的

np.trace

。你可以利用外积的痕迹实际上是内积的事实，这两者正是

np.einsum

的目的。这消除了大量的乘法并且需要更少的内存。

你也可以在你的

np.einsum

中总结最终维度并将输出除以最终维度的长度以摆脱你的

.mean(axis=3)

。

为了在一个

np.einsum

步骤中完成所有操作，您必须首先重塑其中一个输入数组。最终归结为：

A = A.reshape([4, 8, 4, 8, len(t_list), Mmax])

X = np.einsum('ikok...q,ll...q->io...', A, B).reshape((4, 4, len(t_list))) / Mmax

X

现在保存

for

循环的堆叠结果，即最终轴上的所有

ptrace

结果。