Question

我正在尝试执行以下操作

import numpy as np
arr = np.array([0,1,2,3])
fn = lambda x: np.array([x/2,x*2])
res = np.apply_along_axis(fn,0,arr)

也就是说，我希望np.apply_along_axis使用一个返回数组而不是单个值的函数。但是，这会导致以下错误。

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/lib/python3.5/site-packages/numpy/lib/shape_base.py", line 117, in apply_along_axis
    outarr[tuple(i.tolist())] = res
ValueError: could not broadcast input array from shape (2,4) into shape (2)

虽然我希望结果等同于：

res = np.array([[0.0,0],[0.5,2],[1.0,4],[1.5,6]])

请注意，我的实际用例比给定的fn更复杂。我想知道是否有一种方法可以将np.apply_along_axis与返回数组的函数一起使用，或者是否有不同的（numpythonic）方法来实现同样的事情。

谢谢！

Answer 1

这种事情 - 在数组中应用常规Python函数 - 基本上与for循环相同。即使您使用NumPy函数调用您的函数，它也不会非常快。所以请保持简单：

arr = np.array([0,1,2,3])
fn = lambda x: [x/2,x*2]
res = np.array(list(map(fn, arr)))

或者如果你关心表现：

arr = np.array([0,1,2,3])
funcs = [lambda x: x/2.0, lambda x: x*2.0]
res = np.empty((len(arr), len(funcs)))
for i, func in enumerate(funcs):
    res[:,i] = func(arr)

第二种方法要快得多，只要数组中的行数多于要应用的函数数。这是因为你只循环几次，使用NumPy矢量化代码一次将每个函数应用于整个输入，而不是逐个循环输入行。

Answer 2

我找到的解决方法是堆叠结果，使lambda返回一维数组而不是二维数组：

arr = np.array([0, 1, 2, 3])
fn = lambda x: np.hstack([x/2, x*2])
res = np.apply_along_axis(fn, 0, arr)

然后您可以根据需要重新塑造res

Python：使用带有函数返回数组的np.apply_along_axis

2 个答案: