使用 numpy 向量化矩陣運算

Question

我在下面編寫了使用 for 循環的代碼。 我想問是否有辦法在第二個 for 循環中對操作進行矢量化，因為我打算使用更大的矩陣。

import numpy as np

num = 5

A = np.array([[1,2,3,4,5], [4,5,6,4,5], [7,8,9,4,5], [10,11,12,4,5], [13,14,15,4,5]])
sm_factor = np.array([0.1 ,0.1, 0.1, 0.1, 0.1])


d2m = np.zeros((num, num))
d2m[0, 0] = 2
d2m[0, 1] = -5
d2m[0, 2] = 4
d2m[0, 3] = -1

for k in range(1, num-1):
    d2m[k, k-1] = 1
    d2m[k, k] = -2
    d2m[k, k+1] = 1

d2m[num-1, num-4] = -1
d2m[num-1, num-3] = 4
d2m[num-1, num-2] = -5
d2m[num-1, num-1] = 2

x_smf  = 0
for i in range(len(sm_factor)):
    x_smf = x_smf + sm_factor[i] * (d2m @ (A[i, :]).T).T @ (d2m @ (A[i, :]).T)

x_smf
# 324.0

Answer 1

在對循環的中間結果進行一些研究和打印輸出之后，我找到了解決方案：

x_smf = (sm_factor * np.square(d2m @ A.T).sum(axis=0)).sum()

結果是：

324.0

順便說一句： dm2的創建可以縮短為：

d2m = np.zeros((num, num), dtype='int')
d2m[0, :4] = [ 2, -5,  4, -1]
for k in range(1, num-1):
    d2m[k, k-1:k+2] = [ 1, -2,  1]
d2m[num-1, 1:] = [-1, 4, -5, 2]

Answer 2

您可以使用sps.diags創建稀疏三對角矩陣來避免d2m矩陣創建和x_smf矢量計算的循環，您可以將其轉換為數組以便能夠編輯第一行和最后一行。 您的代碼將如下所示（請注意，第 10 行中的診斷結果已使用diags方法轉換為密集的ndarray ）：

import numpy as np
import scipy.sparse as sps

# Dense tridiagonal matrix
d2m = sps.diags([1, -2, 1], [-1, 0, 1], shape=(num, num)).toarray() # cast to array
# First line boundary conditions
d2m[0, 0] = 2
d2m[0, 1] = -5
d2m[0, 2] = 4
d2m[0, 3] = -1
# Last line boundary conditions
d2m[num-1, num-4] = -1
d2m[num-1, num-3] = 4
d2m[num-1, num-2] = -5
d2m[num-1, num-1] = 2

Valdi_Bo提出的解決方案使您能夠刪除第二個 FOR 循環：

x_smf = np.sum(sm_factor * np.square(d2m @ A.T).sum(axis=0))

但是，我想引起您的注意， x_smf矩陣是稀疏的，並且將其存儲為密集的 ndarray 對計算時間和 memory 存儲都是不利的。 我建議您不要轉換為密集的 ndarray，而是轉換為稀疏矩陣格式。 例如lil_matrix ，它是列表稀疏矩陣格式的列表，使用 tolil() 方法而不是 toarray()：

# Sparse tridiagonal matrix
d2m_s = sps.diags([1, -2, 1], [-1, 0, 1], shape=(num, num)).tolil() # cast to lil

這是一個腳本，它比較了更大情況下的三個實現num=4000 （對於num=5都給出324 ）。 對於這種大小，我已經看到了使用稀疏矩陣的好處，這是整個腳本（第一行是num與5不同的代碼的概括）：

from time import time
import numpy as np
import scipy.sparse as sps

num = 4000

A = np.concatenate([np.arange(1, (num-2)*num+1).reshape(num, num-2), np.repeat([[4, 5]], num, axis=0)], axis=1)
sm_factor = 0.1*np.ones(num)

########## DENSE matrix + FOR loop ##########

d2m = sps.diags([1, -2, 1], [-1, 0, 1], shape=(num, num)).toarray() # cast to array
# First line boundary conditions
d2m[0, 0] = 2
d2m[0, 1] = -5
d2m[0, 2] = 4
d2m[0, 3] = -1
# Last line boundary conditions
d2m[num-1, num-4] = -1
d2m[num-1, num-3] = 4
d2m[num-1, num-2] = -5
d2m[num-1, num-1] = 2

# FOR loop version
t_start = time()
x_smf  = 0
for i in range(len(sm_factor)):
    x_smf = x_smf + sm_factor[i] * (d2m @ (A[i, :]).T).T @ (d2m @ (A[i, :]).T)
print(f'FOR loop version time:           {time()-t_start}s')
print(f'FOR loop version value:          {x_smf}\n')

########## DENSE matrix + VECTORIZED ##########

t_start = time()
x_smf_v = np.sum(sm_factor * np.square(d2m @ A.T).sum(axis=0))
print(f'VECTORIZED version time:         {time()-t_start}s')
print(f'VECTORIZED version value:        {x_smf_v}\n')

########## SPARSE matrix + VECTORIZED ##########

d2m_s = sps.diags([1, -2, 1], [-1, 0, 1], shape=(num, num)).tolil() # cast to lil
# First line boundary conditions
d2m_s[0, 0] = 2
d2m_s[0, 1] = -5
d2m_s[0, 2] = 4
d2m_s[0, 3] = -1
# Last line boundary conditions
d2m_s[num-1, num-4] = -1
d2m_s[num-1, num-3] = 4
d2m_s[num-1, num-2] = -5
d2m_s[num-1, num-1] = 2

t_start = time()
x_smf_s = np.sum(sm_factor * np.square(d2m_s @ A.T).sum(axis=0))
print(f'SPARSE+VECTORIZED version time:  {time()-t_start}s')
print(f'SPARSE+VECTORIZED version value: {x_smf_s}\n')

這是我在運行代碼時得到的：

FOR loop version time:           25.878241777420044s
FOR loop version value:          3.752317536763356e+17

VECTORIZED version time:         1.0873610973358154s
VECTORIZED version value:        3.752317536763356e+17

SPARSE+VECTORIZED version time:  0.37279224395751953s
SPARSE+VECTORIZED version value: 3.752317536763356e+17

正如您所看到的，使用稀疏矩陣可以讓您在計算時間上贏得另一個因素 3，並且不需要您調整之后的代碼。 測試稀疏矩陣的各種 scipy 實現（tocsc()、tocsr()、todok() 等）也是一個很好的策略，有些可能更適合您的情況。