如何使用 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, :, :]