创建组合的 numpy 数组，无需显式循环/理解

Question

我有一个很长的 NumPy 数组，我需要将其组织为“模式”，如下所示：

import numpy as np

# long array:
a = np.array(
    [
        1.3,
        -1.8,
        0.3,
        11.4,
        # ...
    ]
)


def pattern(x: float):
    return np.array(
        [
            [x, 0, 0],
            [0, x, 0],
            [0, 0, x],
            [+x, +x, +x],
            [-x, -x, -x],
        ]
    )


out = np.array([pattern(x) for x in a])
print(out.shape)

(4, 5, 3)

我想知道是否有一种方法可以在out 不显式循环的情况下进行构造a。

有任何想法吗？

Answer 1

将函数替换pattern为 -1/0/1 值的数组并广播乘法：

\n

pattern = np.array([[ 1,  0,  0],\n                    [ 0,  1,  0],\n                    [ 0,  0,  1],\n                    [+1, +1, +1],\n                    [-1, -1, -1]])\n\nout = a[:,None,None] * pattern[None]\n

\n

变体einsum：

\n

out = np.einsum(\'i,jk->ijk\', a, pattern)\n

\n

另一种方法select（效率较低）：

\n

m1 = np.array([[ True, False, False],\n               [False,  True, False],\n               [False, False,  True],\n               [ True,  True,  True],\n               [False, False, False]])\nm2 = np.array([[False, False, False],\n               [False, False, False],\n               [False, False, False],\n               [False, False, False],\n               [ True,  True,  True]])\n\nx = a[:,None,None]\nout = np.select([m1, m2], [x, -x], 0)\n

\n

输出：

\n

array([[[  1.3,   0. ,   0. ],\n        [  0. ,   1.3,   0. ],\n        [  0. ,   0. ,   1.3],\n        [  1.3,   1.3,   1.3],\n        [ -1.3,  -1.3,  -1.3]],\n\n       [[ -1.8,  -0. ,  -0. ],\n        [ -0. ,  -1.8,  -0. ],\n        [ -0. ,  -0. ,  -1.8],\n        [ -1.8,  -1.8,  -1.8],\n        [  1.8,   1.8,   1.8]],\n\n       [[  0.3,   0. ,   0. ],\n        [  0. ,   0.3,   0. ],\n        [  0. ,   0. ,   0.3],\n        [  0.3,   0.3,   0.3],\n        [ -0.3,  -0.3,  -0.3]],\n\n       [[ 11.4,   0. ,   0. ],\n        [  0. ,  11.4,   0. ],\n        [  0. ,   0. ,  11.4],\n        [ 11.4,  11.4,  11.4],\n        [-11.4, -11.4, -11.4]]])\n

\n

时间：

\n

# list comprehension\n14.8 \xc2\xb5s \xc2\xb1 1.26 \xc2\xb5s per loop (mean \xc2\xb1 std. dev. of 7 runs, 10,000 loops each)\n\n# broadcasting\n1.87 \xc2\xb5s \xc2\xb1 85.4 ns per loop (mean \xc2\xb1 std. dev. of 7 runs, 1,000,000 loops each)\n\n# einsum\n3.07 \xc2\xb5s \xc2\xb1 184 ns per loop (mean \xc2\xb1 std. dev. of 7 runs, 100,000 loops each)\n\n# select\n42.2 \xc2\xb5s \xc2\xb1 2.32 \xc2\xb5s per loop (mean \xc2\xb1 std. dev. of 7 runs, 10,000 loops each)\n

\n