沿numpy数组的任意轴迭代连续的一维切片答案

【问题标题】：Iterating over successive 1-D slices along arbitrary axis of numpy array沿numpy数组的任意轴迭代连续的一维切片
【发布时间】：2019-09-18 11:58:01
【问题描述】：

我正在编写一个 python 包，它沿着任意形状的numpy 数组的任意轴执行各种复杂的统计分析任务。

目前，为了使数组的形状和轴可以是任意的，我只是对数组进行置换，使感兴趣的轴放置在远 RHS 上，并将 LHS 轴压缩为一个。例如，如果数组形状是(3,4,5)，我们想沿轴1进行一些操作，则将其转换为形状(15,4)，沿轴进行操作-1，然后将其转换回进入形状(3,4,5) 并由函数返回。

由于所有这些数组操作，我觉得这种方法可能会不必要地缓慢。有没有一种方法可以干净地迭代数组的除了一个维度之外的所有维度？也就是说，在上面的示例中，这将变为 [0,:,0]、[0,:,1]、...、[2,:,3]、[2,:,4]，但这同样适用于任意数组形状和轴位置。

也许np.ndenumerate、np.ndindex 和np.take 可以用于此目的？

编辑：有没有办法用np.nditer 做到这一点？也许这可以匹配排列/重塑的速度。

【问题讨论】：

np.apply_along_axis 可能对你有用。但是看看它的代码。它将所需的轴转置到末尾，使用ndindex 生成索引，并对它们进行迭代，然后将结果转置回来。听起来很像你的方法。
您的置换和重塑并不那么昂贵。在某些情况下，它只是一个视图，但通常它会生成一个副本。但与可能便宜的迭代相比。但是做一些时间安排；不要只凭感觉。 apply... docs 声称比使用拆分 ndindex 的同等速度更快，但我不知道差异是否显着。我建议不要在不需要额外行李的简单二维情况下使用它。
您能否为我们提供一些可以使用的示例数组？特别是您希望在此轴上应用的功能。就内存或时间而言，轴置换不是一项昂贵的操作。所以，我们必须定时看看@hpaulj 提到的......
@kmario23 尝试了一下，你们是对的，排列和重塑确实快得多（见下面的答案）。这真让我吃惊。本来以为在某些情况下它会更慢，即当 RAM 成为瓶颈时......但情况似乎并非如此。
@LukeDavis 如果您真的关心内存并希望避免在重塑期间复制，请查看Notes in Numpy reshape。它提供了一种在发生数据复制时引发错误的机制。

标签： python arrays numpy reshape numpy-ndarray

【解决方案1】：

事实证明，只是转置和重塑确实更快。所以我想答案是......不要那样做，最好像我已经在做的那样置换和重塑。

这是我项目中的代码。

# Benchmark
f = lambda x: x # can change this to any arbitrary function
def test1(data, axis=-1):
    # Test the lead flatten approach
    data, shape = lead_flatten(permute(data, axis))
    output = np.empty(data.shape)
    for i in range(data.shape[0]): # iterate along first dimension; each row is an autocor
        output[i,:] = f(data[i,:]) # arbitrary complex equation
    return unpermute(lead_unflatten(output, shape), axis)
def test2(data, axis=-1):
    # Test the new approach
    output = np.empty(data.shape)
    for d,o in zip(iter_1d(data, axis), iter_1d(output, axis)):
        o[...] = f(d)
    return output

# Iterator class
class iter_1d(object):
    def __init__(self, data, axis=-1):
        axis = (axis % data.ndim) # e.g. for 3D array, -1 becomes 2
        self.data = data
        self.axis = axis
    def __iter__(self):
        shape = (s for i,s in enumerate(self.data.shape) if i!=self.axis)
        self.iter = np.ndindex(*shape)
        return self
    def __next__(self):
        idx = self.iter.next()
        idx = [*idx]
        idx.insert(self.axis, slice(None))
        return self.data[idx]

# Permute and reshape functions
def lead_flatten(data, nflat=None):
    shape = list(data.shape)
    if nflat is None:
        nflat = data.ndim-1 # all but last dimension
    if nflat<=0: # just apply singleton dimension
        return data[None,...], shape
    return np.reshape(data, (np.prod(data.shape[:nflat]).astype(int), *data.shape[nflat:]), order='C'), shape # make column major

def lead_unflatten(data, shape, nflat=None):
    if nflat is None:
        nflat = len(shape) - 1 # all but last dimension
    if nflat<=0: # we artificially added a singleton dimension; remove it
        return data[0,...]
    if data.shape[0] != np.prod(shape[:nflat]):
        raise ValueError(f'Number of leading elements {data.shape[0]} does not match leading shape {shape[nflat:]}.')
    if not all(s1==s2 for s1,s2 in zip(data.shape[1:], shape[nflat:])):
        raise ValueError(f'Trailing dimensions on data, {data.shape[1:]}, do not match trailing dimensions on new shape, {shape[nflat:]}.')
    return np.reshape(data, shape, order='C')

def permute(data, source=-1, destination=-1):
    data = np.moveaxis(data, source, destination)
    return data

def unpermute(data, source=-1, destination=-1):
    data = np.moveaxis(data, destination, source)
    return data

下面是一些%timeit 操作的结果。

import numpy as np
a = np.random.rand(10,20,30,40)
%timeit -r10 -n10 test1(a, axis=2) # around 12ms
%timeit -r10 -n10 test2(a, axis=2) # around 22ms

【讨论】：