向量化 numpy 操作

Question

我有一个一维 numpy ndarray V。我的操作是 H = (V/a)**b，其中 a 和 b 只是标量。 这很简单，但我需要运行多个实验（数千个），在其中尝试 a 和 b 的不同变体。 使用 for 循环很容易做到这一点，但我希望它尽可能快地运行，所以我想矢量化。 所以，假设我也生成了 a 和 b 的一维数组。 给定 3 个一维数组 V、a 和 b，我该如何有效地做到这一点？ 谢谢！

Answer 1

你可以这样做：

import numpy as np

V = np.arange(100.0)
a = np.random.rand(50)
b = np.random.rand(50)

result = (V[np.newaxis, :] / a[:, np.newaxis]) ** b[:, np.newaxis]

result大小为 50x100，因此result[i]将是使用参数a[i]和b[i] 。

Answer 2

这就是广播的用武之地，您可以添加维度，以便数组沿每个轴广播：

(V/a[:, None])**b[:, None, None]

这将向a添加一个尾随维度，向b添加两个尾随维度，并且广播将确保您将获得一个 3D 数组。 这是有效的，因为广播将始终沿着最后一个轴“广播”。

例如：

>>> import numpy as np
>>> V = np.array([1,2,3])
>>> a = np.array([0.1, 0.2, 0.3])
>>> b = np.array([1, 0.5, 0.34])

>>> a[:, None]
array([[ 0.1],
       [ 0.2],
       [ 0.3]])

>>> b[:, None, None]
array([[[ 1.  ]],

       [[ 0.5 ]],

       [[ 0.34]]])

>>> (V/a[:, None])**b[:, None, None]
array([[[ 10.        ,  20.        ,  30.        ],
        [  5.        ,  10.        ,  15.        ],
        [  3.33333333,   6.66666667,  10.        ]],

       [[  3.16227766,   4.47213595,   5.47722558],
        [  2.23606798,   3.16227766,   3.87298335],
        [  1.82574186,   2.5819889 ,   3.16227766]],

       [[  2.18776162,   2.7691737 ,   3.17849276],
        [  1.72842206,   2.18776162,   2.51114059],
        [  1.50583981,   1.90602666,   2.18776162]]])

请注意，这可能会占用大量内存。 因此，如果您的所有数组都很长，则可能无法使用它。

例如，如果所有数组都包含 1000 个值，则结果将包含 10 亿个元素，这会占用大量 RAM（约 7.5GB 与float64 s）。