python - numpy:沿第一轴乘以任意形状数组

Question

我想将一个数组沿其第一个轴乘以某个向量。

例如，如果 a 是二维的，b 是一维的，并且 a.shape[0] == b.shape[0]，我们可以这样做:

a *= b[:, np.newaxis]

如果 a 具有任意形状怎么办？在 numpy 中，省略号“...”可以解释为“用':'填充剩余的索引”。是否有用 None/np.newaxis 填充剩余轴的等价物？

下面的代码生成了所需的结果，但我更喜欢一种通用的矢量化方法来完成此操作，而无需退回到 for 循环。

from __future__ import print_function

import numpy as np


def foo(a, b):
    """
    Multiply a along its first axis by b
    """
    if len(a.shape) == 1:
        a *= b
    elif len(a.shape) == 2:
        a *= b[:, np.newaxis]
    elif len(a.shape) == 3:
        a *= b[:, np.newaxis, np.newaxis]
    else:
        n = a.shape[0]
        for i in range(n):
           a[i, ...] *= b[i]          

n = 10
b = np.arange(n)

a = np.ones((n, 3))
foo(a, b)
print(a)

a = np.ones((n, 3, 3))
foo(a, b)
print(a)

Answer 1

只需颠倒坐标轴的顺序即可:

transpose = a.T
transpose *= b

a.T 是a 的转置 View ，其中“转置”意味着反转任意维a 的维度顺序。我们将 a.T 分配给一个单独的变量，这样 *= 就不会尝试设置 a.T 属性；结果仍然适用于 a，因为转置是一个 View 。

演示:

In [55]: a = numpy.ones((2, 2, 3))

In [56]: a
Out[56]: 
array([[[1., 1., 1.],
        [1., 1., 1.]],

       [[1., 1., 1.],
        [1., 1., 1.]]])

In [57]: transpose = a.T

In [58]: transpose *= [2, 3]

In [59]: a
Out[59]: 
array([[[2., 2., 2.],
        [2., 2., 2.]],

       [[3., 3., 3.],
        [3., 3., 3.]]])