我希望将数组沿着它的第一个轴乘以某个向量。
例如,如果a是2D,b是1D,而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)
答案 0 :(得分:2)
只需颠倒轴的顺序:
transpose = a.T
transpose *= b
a.T
是a
的转置视图,其中"转置"意味着颠倒任意维a
的维度顺序。我们将a.T
分配给单独的变量,以便*=
不会尝试设置a.T
属性;结果仍适用于a
,因为转置是一种观点。
演示:
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.]]])