使用 Numpy 向量化三个嵌套循环

Question

我有一个复数矩阵 C，尺寸为 (r, r)，还有一个复数向量 r。我需要根据以下等式从 C 和 v 计算一个新矩阵：

其中 K 也是维度为 (r, r) 的方阵。下面是使用三个循环计算 K 的代码：

import numpy as np
import matplotlib.pyplot as plt

r = 9

#   Create random matrix
C = np.random.rand(r,r) + np.random.rand(r,r) * 1j
v = np.random.rand(r) + np.random.rand(r) * 1j

#   Original loops
K = np.zeros((r, r))

for m in range(r):
    for n in range(r):
        for i in range(r):
            K[m,n] += np.imag( C[i,m] * np.conj(C[i,n]) * np.sign(np.imag(v[i])) )

plt.figure()
plt.imshow(K)
plt.show()

使用 i 删除循环相对容易：

#   First optimization
K = np.zeros((r, r))

for m in range(r):
    for n in range(r):
        K[m,n] = np.imag(np.sum(C[:,m] * np.conj(C[:,n]) * np.sign(np.imag(v)) ))

但我不确定如何继续向量化剩余的两个循环。在这种情况下真的可以吗？

Answer 1

我找到了一些暂时可用的东西。但是，仍然存在一个循环，并且由于生成的矩阵是对称的，因此仍然需要进行一些优化。

我没有删除 i 循环，而是删除了另外两个循环：

K = np.zeros((r, r), dtype=np.complex128)

for i in range(r):
    K += adjointMatrix(C) @ (np.sign(np.imag(v)) * C)

K = np.imag(K)

与：

def adjointMatrix(X):
    return np.conjugate( np.transpose(X) )

Answer 2

这看起来像矩阵乘法：

out = np.imag((C*np.sign(np.imag(v))[:,None]).T @ np.conj(C))

或者您可以使用 np.einsum:

out = np.imag(np.einsum('im,in,i', C, np.conj(C), np.sign(np.imag(v))))

验证你的方法：

np.all(np.abs(out-K) < 1e-6)
# True

Answer 3

我有很多这样的问题，下面是我通常如何着手寻找解决方案来编写矢量化代码。

以下是我注意到的关于您的总结。很酷的结论是，您可能根本不需要矢量化，因为您可以将整个计算表示为二维矩阵的单一乘积。来了...

让我们首先定义以下矩阵（抱歉缺少 Latex 符号，Whosebug 不支持 Mathjax）：

A_{i,j} = c_{i,j}.

B_{i,j} = c_{i,j} * sgn(Im(v_i))

那么你可以把你的总和写成：

k_{m,n} = Im( \sum_{i=1}^{r} c_{i,m} * sgn(Im(v_i)) * c_{i,n} ^* ) = Im ( \sum_{i=1}^{r} B_{i,m} * A_{i,n}^* ) = Im( \sum_{i=1}^{r} B_{m ,i}^T * A_{i,n}^* )

上面 Im(.) 内部的表达式是矩阵乘法的定义，等效于以下内容：

k_{m,n} = Im( (B^T * A^*)_{m,n} )

这意味着您的矩阵 k 可以表示为矩阵 B 的转置与矩阵 A 的乘积。在您的代码中，矩阵 A 已分配给变量 C。因此可以按如下方式进行矢量化：

C = np.random.rand(r,r) + np.random.rand(r,r) * 1j
v = np.random.rand(r) + np.random.rand(r) * 1j
k = np.imag( (C * np.sign(np.imag(v)).T @ np.conj(C) )

并且您避免了讨厌的循环和令人费解的表达式

使用 Numpy 向量化三个嵌套循环

Vectorizing three nested loops with Numpy

python

loops

numpy

vectorization