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

Question

我最近阅读了很多关于迭代 numpy 数组的不同技术，似乎共识是根本不迭代(例如，参见 a comment here )。 SO 上有几个类似的问题，但我的情况有点不同，因为我必须结合“迭代”(或不迭代)和访问以前的值。

假设有 N(N 很小，通常为 4，可能最多为 7)个 float128 的一维 numpy 数组在列表中 X ，所有数组的大小相同。为了给你一点洞察力，这些是来自 PDE 积分的数据，每个数组代表一个函数，我想应用一个 Poincare 截面。不幸的是，该算法应该既节省内存又节省时间，因为这些数组有时每个大约 1Gb，并且板上只有 4Gb 的 RAM(我刚刚了解了 numpy 数组的内存映射，现在考虑改用它们常规的)。

其中一个数组用于“过滤”其他数组，因此我从 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) )

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

另一种方法是使用 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 是)。读取一个数组，然后为一个 2-D numpy 数组分配内存(这里的内存开销很小)并将剩余的读取到该 2-D 数组中会更有效吗？

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