如何使用 NumPy 通過 2D 數組以矢量化方式縮放一組 2D arrays（3D 數組）？

Question

我有一個 3D 矩陣，其中包含 M 個通道 [M x N x N] 的 N x N 個協方差矩陣。 我還有一個在一系列時間點 [M x T] 的每個通道的縮放因子的 2D 矩陣。 我想生成一個 4D 矩陣，其中包含每個時間點相關通道協方差的縮放版本。 所以要清楚，[M x T] * [M x N x N] -> [M x T x N x N]

當前版本使用 for 循環：

m, t, n = 4, 10, 7
channel_timeseries = np.zeros((m, t))
covariances = np.random.rand(m, n, n)

result_array = np.zeros((m, t, n, n))

# Each channel
for i, (channel_cov, channel_timeseries) in enumerate(zip(covariances, channel_timeseries)):
    # Each time point
    for j, time_point in enumerate(channel_timeseries):
        result_array[i, j] = time_point * channel_cov

這應該導致結果數組全為零。 用np.ones替換 channel_timeseries 的初始化，我們應該看到每個通道的協方差在時間序列的每一步都沒有改變。

對我來說真正重要的情況是每個通道在每個時間點都有一個標量值，我們通過匹配正確通道和時間點的值來縮放相關通道的協方差矩陣。

正如你在上面看到的，我可以用一個 for 循環來做到這一點，它工作得很好，但我正在處理一些巨大的數據集，最好有一個矢量化的解決方案。

非常感謝您的時間。

Answer 1

numpy.einsum將在這里派上用場。 我用隨機channel_timeseries數組修改了你的代碼，增加了 arrays 的大小，並重命名了循環變量（否則你會覆蓋原來的變量！）

import numpy as np
import time

m, t, n = 40, 100, 70
channel_timeseries = np.random.rand(m, t)
covariances = np.random.rand(m, n, n)

t0 = time.time()
result_array_1 = np.zeros((m, t, n, n))
# Each channel
for i, (c_cov, c_ts) in enumerate(zip(covariances, channel_timeseries)):
    # Each time point
    for j, time_point in enumerate(c_ts):
        result_array_1[i, j] = time_point * c_cov
t1 = time.time()
result_array_2 = np.einsum('ij,ikl->ijkl', channel_timeseries, covariances)
t2 = time.time()

print(np.array_equal(result_array_1, result_array_2)) # True
print('Time for result_array_1: ', t1-t0) # 0.07601261138916016
print('Time for result_array_2: ', t2-t1) # 0.02957916259765625

這導致我的機器中的numpy.einsum的速度提高了 50% 以上。

Answer 2

正如 b-fg 所說，您可以使用np.einsum

np.einsum('mt,mno->mtno', channel_timeseries, covariances)

或廣播：

channel_timeseries[:, :, None, None] * covariances[:, None, :, :]