在numpy.sum()或mean()之前有效的numpy.roll

Question

我有几个(1000个订单)3D形状阵列(1000,800,1024)我想学习.我需要计算沿轴= 0的平均值,但在我能做到之前,我必须沿着轴2滚动数据,直到它"在正确的位置".

这听起来很奇怪,所以我会试着解释一下.形状(1024,)的1D子阵列是来自物理环形缓冲器的数据.我知道,环形缓冲区以不同的方式读出.所以我有几个pos形状阵列(1000,800).告诉我环读缓冲区被读出的位置.我的3D data形状阵列(1000,800,1024)我需要根据它滚动pos.

只有在滚动之后...... 3D阵列对我来说才有意义,我可以开始分析它们.在C可以写出这是做这个非常简单的代码,所以我不知道如果我可以种在的结束"告诉"的numpy的平均值()或SUM(),他们应该在不同的指数开始例程和"打滚" 1D-子阵列.

我目前做的是这样的:

rolled = np.zeros_like(data) # shape (1000, 800, 1024)
for a in range(rolled.shape[0]):
    for b in range(rolled.shape[1]):
        rolled[a,b] = np.roll(data[a,b], pos[a,b])

这需要大约60秒然后我做,例如:

m = rolled.mean(axis=0)
s = rolled.std(axis=0)

这只需要15秒左右.

我的观点是,使轧制副本需要花费很多的时间和空间(好吧,我可以通过写卷的东西放回节省空间data),而肯定是有办法(C语言)来实现这种平均在一个循环滚动因此节省了大量时间.我的问题是......如果有办法与numpy做类似的事情？

Answer 1

我感到无聊,并在Cython中编写了你的​​函数.它运行速度比您发布的代码快10倍,而无需分配中间数组.

import numpy as np
cimport numpy as np
cimport cython
from libc.math cimport sqrt

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.nonecheck(False)
@cython.cdivision(True)
def rolled_mean_std(double[:,:,::1] data, int[:,::1] pos):
  cdef Py_ssize_t s1,s2,s3,a,b,c
  s1 = data.shape[0]
  s2 = data.shape[1]
  s3 = data.shape[2]
  cdef double[:,::1] sums = np.zeros((s2,s3))
  cdef double[:,::1] sumsq = np.zeros((s2,s3))
  cdef double d
  cdef int p
  # Compute sums and sum-of-squares.
  for a in range(s1):
    for b in range(s2):
      p = pos[a,b]
      for c in range(s3):
        d = data[a,b,(c+s3-p)%s3]
        sums[b,c] += d
        sumsq[b,c] += d * d
  # Calculate mean + std in place.
  for b in range(s2):
    for c in range(s3):
      d = sums[b,c]
      sums[b,c] /= s1
      sumsq[b,c] = sqrt((s1*sumsq[b,c] - (d*d)))/s1
  return sums, sumsq

请注意,这使用了Naive mean + stdv算法,因此您可能会遇到浮点精度错误.不过,我的测试没有显示出太大的影响.