迭代几个 numpy 数组并处理当前和先前元素的有效方法？

Question

我最近阅读了很多关于迭代 numpy 数组的不同技术的文章，似乎一致认为根本不迭代（例如，请参阅 a comment here）。关于 SO 有几个类似的问题，但我的情况有点不同，因为我必须结合“迭代”（或不迭代）和访问以前的值。

假设列表 X 中有 N 个（N 很小，通常为 4，可能最多 7 个）float128 的一维 numpy 数组，所有数组的大小相同。为了让您了解一下，这些是 PDE 集成的数据，每个数组代表一个函数，我想应用一个 Poincare 部分。不幸的是，该算法应该既节省内存又节省时间，因为这些数组有时每个约为 1Gb，并且板上只有 4Gb 的 RAM（我刚刚了解了 numpy 数组的 memmap'ing，现在考虑使用它们常规的）。

其中一个数组用于“过滤”其他数组，所以我从secaxis = X.pop(idx) 开始。现在我必须找到(secaxis[i-1] > 0 and secaxis[i] < 0) or (secaxis[i-1] < 0 and secaxis[i] > 0) 的索引对，然后对剩余的数组X 应用简单的代数变换（并保存结果）。值得一提的是，在此操作期间不应浪费数据。

有多种方法可以做到这一点，但对我来说，没有一种方法看起来很有效（也不够优雅）。一种是类似 C 的方法，您只需在 for 循环中进行迭代：

import array # better than lists
res = [ array.array('d') for _ in X ]
for i in xrange(1,secaxis.size):
  if condition: # see above
    co = -secaxis[i-1]/secaxis[i]
    for j in xrange(N):
      res[j].append( (X[j][i-1] + co*X[j][i])/(1+co) )

这显然是非常低效的，而且不是 Python 的方式。

另一种方法是使用 numpy.nditer，但我还没有弄清楚如何访问前一个值，尽管它允许一次迭代多个数组：

# without secaxis = X.pop(idx)
it = numpy.nditer(X)
for vec in it:
  # vec[idx] is current value, how do you get the previous (or next) one?

第三种可能性是首先找到具有高效 numpy 切片的搜索索引，然后将它们用于批量乘法/加法。我现在更喜欢这个：

res = []
inds, = numpy.where((secaxis[:-1] < 0) * (secaxis[1:] > 0) +
                   (secaxis[:-1] > 0) * (secaxis[1:] < 0))
coefs = -secaxis[inds] / secaxis[inds+1] # array of coefficients
for f in X: # loop is done only N-1 times, that is, 3 to 6
    res.append( (f[inds] + coefs*f[inds+1]) / (1+coefs) )

但这似乎是在 7 + 2*(N - 1) 次传递中完成的，此外，我不确定secaxis[inds] 的寻址类型（它不是切片，通常它必须通过索引查找所有元素就像第一种方法一样，不是吗？）。

最后，我也尝试过使用 itertools，但它导致了巨大而晦涩的结构，这可能源于我对函数式编程不是很熟悉：

def filt(x):
  return (x[0] < 0 and x[1] > 0) or (x[0] > 0 and x[1] < 0)
import array
from itertools import izip, tee, ifilter
res = [ array.array('d') for _ in X ] 
iters = [iter(x) for x in X]   # N-1 iterators in a list
prev, curr = tee(izip(*iters)) # 2 similar iterators, each of which
                               # consists of N-1 iterators
next(curr, None) # one of them is now for current value
seciter = tee(iter(secaxis))
next(seciter[1], None)
for x in ifilter(filt, izip(seciter[0], seciter[1], prev, curr)):
  co = - x[0]/x[1]
  for r, p, c in zip(res, x[2], x[3]):
    r.append( (p+co*c) / (1+co) )

这不仅看起来很丑，而且还需要很长时间才能完成。

所以，我有以下问题：

1. 在所有这些方法中，第三种方法确实是最好的吗？如果是这样，可以做些什么来改善最后一个？
2. 还有其他更好的吗？
3. 出于好奇，有没有办法使用 nditer 解决问题？
4. 最后，使用 numpy 数组的 memmap 版本会更好，还是会减慢速度？也许我应该只将secaxis 数组加载到 RAM 中，将其他数组保留在磁盘上并使用第三种方法？
5. （额外问题）长度相等的一维 numpy 数组列表来自加载 N 个 .npy 文件，其大小事先未知（但 N 是）。读取一个数组，然后为一个二维 numpy 数组分配内存（这里的内存开销很小）并将剩余部分读入该二维数组会更有效吗？

Answer 1

numpy.where() 版本已经足够快了，你可以通过method3() 稍微加快它。如果>条件可以改为>=，也可以使用method4()。

import numpy as np

a = np.random.randn(100000)

def method1(a):
    idx = []
    for i in range(1, len(a)):
        if (a[i-1] > 0 and a[i] < 0) or (a[i-1] < 0 and a[i] > 0):
            idx.append(i)
    return idx

def method2(a):
    inds, = np.where((a[:-1] < 0) * (a[1:] > 0) +
                       (a[:-1] > 0) * (a[1:] < 0))
    return inds + 1

def method3(a):
    m = a < 0
    p = a > 0
    return np.where((m[:-1] & p[1:]) | (p[:-1] & m[1:]))[0] + 1

def method4(a):
    return np.where(np.diff(a >= 0))[0] + 1

assert np.allclose(method1(a), method2(a))
assert np.allclose(method2(a), method3(a))
assert np.allclose(method3(a), method4(a))

%timeit method1(a)
%timeit method2(a)
%timeit method3(a)
%timeit method4(a)

%timeit 结果：

1 loop, best of 3: 294 ms per loop
1000 loops, best of 3: 1.52 ms per loop
1000 loops, best of 3: 1.38 ms per loop
1000 loops, best of 3: 1.39 ms per loop

Answer 2

我需要更详细地阅读您的帖子，但将从一些一般性观察开始（来自以前的迭代问题）。

在 Python 中没有一种有效的迭代数组的方法，尽管有些事情会减慢速度。我喜欢区分迭代机制（nditer、for x in A:）和动作（alist.append(...)、x[i+1] += 1）。大消费者通常是动作，做了很多次，而不是迭代机制本身。

让numpy在编译后的代码中进行迭代是最快的。

 xdiff = x[1:] - x[:-1]

比

快很多

 xdiff = np.zeros(x.shape[0]-1)
 for i in range(x.shape[0]:
     xdiff[i] = x[i+1] - x[i]

np.nditer 并不快。

nditer 推荐作为编译代码中的通用迭代工具。但它的主要价值在于处理广播和协调多个数组（输入/输出）的迭代。而且您需要使用缓冲和 c 之类的代码来从 nditer 获得最佳速度（我会查找最近的 SO 问题）。

https://stackoverflow.com/a/39058906/901925

如果没有研究相关的iteration 教程页面（以cython 示例结尾的页面），请勿使用nditer。

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

仅从经验来看，这种方法将是最快的。是的，它会多次迭代secaxis，但这些都是在编译代码中完成的，并且比 Python 中的任何迭代都要快得多。而for f in X: 迭代只是几次。

res = []
inds, = numpy.where((secaxis[:-1] < 0) * (secaxis[1:] > 0) +
                   (secaxis[:-1] > 0) * (secaxis[1:] < 0))
coefs = -secaxis[inds] / secaxis[inds+1] # array of coefficients
for f in X: 
    res.append( (f[inds] + coefs*f[inds+1]) / (1+coefs) )

@HYRY 探索了使where 步骤更快的替代方法。但正如你所看到的，差异并没有那么大。其他可能的调整

inds1 = inds+1
coefs = -secaxis[inds] / secaxis[inds1]
coefs1 = coefs+1
for f in X:
    res.append(( f[inds] + coefs*f[inds1]) / coefs1)

如果X 是一个数组，那么res 也可以是一个数组。

res = (X[:,inds] + coefs*X[:,inds1])/coefs1

但对于较小的N，我怀疑列表res 也一样好。不需要使数组比必要的大。调整很小，只是试图避免重新计算。

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

np.where 的用法就是 np.nonzero。这实际上对数组进行了两次传递，一次使用np.count_nonzero 来确定它将返回多少个值，并创建返回结构（现在已知长度的数组列表）。第二个循环来填充这些索引。因此，如果它使操作保持简单，那么多次迭代就可以了。