numpy - 函数返回高维数组时numpy apply_along_axis怎么办？

Question

我有一个形状为 (10, 100000) 的数组和一个函数 f，它采用形状为 (100000,) 的数组来形状为 (200,200) 的数组。将 f 应用于 10 行中的每一行以获得形状数组的最简单方法是什么(10,200,200)？我希望使用 apply_along_axis，但它似乎行不通，因为 f 的输出维数高于输入维数。

(更一般地，给定一个形状为 (x1,...,xn,y1,...,ym) 的数组 a 和一个函数 f 将形状为 (y1,...,ym) 的数组转换为形状为 (z1,...,zp) 的数组，您可能需要将 f 应用于 a 的最后 m 维度，对于 的前 n 维度的每个设置>a，得到形状为 (x1,...,xn,z1,...,zp) 的数组。或者你可能有一个问题可以转置为一个这种形式。进行此类转换的最佳方法是什么？)

Answer 1

我的第一个想法是 reshape a , 折叠第一个 n尺寸缩小到一个。然后这只是迭代该维度的问题，应用 f到每个子数组。将结果收集在列表(或适当大小的数组)中。最后 reshape 。

正如您所描述的那样 x1...xn尺寸只是随行。

查看 apply_along_axis 的代码.它遍历除传递给函数的轴之外的所有轴。它没有做任何你不能用你自己的迭代做的事情。它将处理 x1...xn 上的迭代, 但需要 y维度要折叠到 1，并且需要一个返回相同形状的函数。

那个函数的核心是

res = func1d(arr[tuple(i.tolist())], *args, **kwargs)
outarr[tuple(ind)] = res

哪里outarr已被初始化为正确的大小，并且ind跨过所有维度(一个维度除外)。它有一个 slice res 所在的对象去。

=====================

一个以二维输入数组开始的简单示例:

In [933]: def foo(arr):
     ...:     return arr.reshape(2,-1)
     ...: 
In [934]: source=np.arange(12).reshape(3,4)
In [935]: dest=np.zeros((source.shape[0],2,2),source.dtype)
In [936]: for i,r in enumerate(source):
     ...:     dest[i,...] = foo(r)
     ...:     
In [937]: dest
Out[937]: 
array([[[ 0,  1],
        [ 2,  3]],

       [[ 4,  5],
        [ 6,  7]],

       [[ 8,  9],
        [10, 11]]])

所以这会迭代源的行，生成新数组，并将它们插入到目标的正确位置。要设置目的地，我必须知道什么 foo产生(维度)。

列表追加方法不需要那么多知识:

In [938]: dest=[]
In [939]: for i,r in enumerate(source):
     ...:     dest.append(foo(r))    
In [940]: dest
Out[940]: 
[array([[0, 1],
        [2, 3]]), array([[4, 5],
        [6, 7]]), array([[ 8,  9],
        [10, 11]])]
In [941]: np.array(dest)
...

归结为老问题，“我如何从函数生成新数组？”