performance - numpy数组子维度上的python操作

Question

许多 numpy 函数提供了使用轴 = 参数在特定轴上操作的选项。我的问题是

这个“沿轴”操作是如何实现的？或者，更直接的问题

如何有效地编写自己的函数提供类似的选项？

我注意到 numpy 提供了一个函数 numpy.apply_along_axis如果基本函数输入是一维数组，这将作为答案。

但是如果我的基函数需要多维输入怎么办？例如。沿着前两个维度 (5,6) 找到形状为 (5,6,2,3,4) 的 np 矩阵 A 的二维移动平均值 B？像一个通用函数 B = f_moving_mean(A, axis=(0,1))

我目前的解决方案是使用 numpy.swapaxes 和 numpy.reshape 来实现这一点。一维移动平均函数的示例代码是:

import pandas as pd
import numpy as np
def nanmoving_mean(data,window,axis=0):
    kw = {'center':True,'window':window,'min_periods':1}
    if len(data.shape)==1:
        return pd.Series(data).rolling(**kw).mean().as_matrix()
    elif len(data.shape)>=2:
        tmp = np.swapaxes(data,0,axis)
        tmpshp = tmp.shape
        tmp = np.reshape( tmp, (tmpshp[0],-1), order='C' )
        tmp = pd.DataFrame(tmp).rolling(**kw).mean().as_matrix()
        tmp = np.reshape( tmp, tmpshp, order='C' )
        return np.swapaxes(tmp,0,axis)
    else:
        print('Invalid dimension!')
        return None

data = np.random.randint(10,size=(2,3,6))
print(data)
nanmoving_mean(data,window=3,axis=2)

这是问题 2 的常见/有效实现方式吗？欢迎任何改进/建议/新方法。

附注。我在这里涉及 pandas 的原因是它的 rolling(...).mean() 方法能够正确处理 nan 数据。

编辑:
我猜问这个问题的另一种方式可能是:循环“动态”维数的语法是什么？

Answer 1

我们可以有一个使用 2D convolution 的方法.

基本步骤是:

作为预处理步骤，替换 NaNs与 0s因为我们需要对输入数据进行加窗求和。

使用 Scipy's convolve2d 获取窗口求和对于数据值
还有NaNs的面具.我们将使用边界元素作为零。

减去 NaNs 的窗口计数从窗口大小获取负责求和的有效元素的计数。

对于边界元素，我们将逐渐减少元素来解释总和。

现在，这些 intervaled-summations也可以通过 Scipy's 1D uniform-filter 获得那相对来说效率更高。另一个好处是我们可以指定要执行这些求和/平均的轴。

混合 Scipy 的 2D convolution和 1D uniform filter ，下面列出的方法很少。

导入相关的 Scipy 函数 -

from scipy.signal import convolve2d as conv2
from scipy.ndimage.filters import uniform_filter1d as uniff

方法#1:

def nanmoving_mean_numpy(data, W): # data: input array, W: Window size
    N = data.shape[-1]
    hW = (W-1)//2

    nan_mask = np.isnan(data)
    data1 = np.where(nan_mask,0,data)

    value_sums = conv2(data1.reshape(-1,N),np.ones((1,W)),'same', boundary='fill')
    nan_sums = conv2(nan_mask.reshape(-1,N),np.ones((1,W)),'same', boundary='fill')

    value_sums.shape = data.shape
    nan_sums.shape = data.shape

    b_sizes = hW+1+np.arange(hW) # Boundary sizes
    count = np.hstack(( b_sizes , W*np.ones(N-2*hW), b_sizes[::-1] ))
    return value_sums/(count - nan_sums)

方法#2:

def nanmoving_mean_numpy_v2(data, W): # data: input array, W: Window size    
    N = data.shape[-1]
    hW = (W-1)//2

    nan_mask = np.isnan(data)
    data1 = np.where(nan_mask,0,data)

    value_sums = uniff(data1,size=W, axis=-1, mode='constant')*W
    nan_sums = conv2(nan_mask.reshape(-1,N),np.ones((1,W)),'same', boundary='fill')
    nan_sums.shape = data.shape

    b_sizes = hW+1+np.arange(hW) # Boundary sizes
    count = np.hstack(( b_sizes , W*np.ones(N-2*hW,dtype=int), b_sizes[::-1] ))
    out =  value_sums/(count - nan_sums)
    out = np.where(np.isclose( count, nan_sums), np.nan, out)
    return out

方法#3:

def nanmoving_mean_numpy_v3(data, W): # data: input array, W: Window size
    N = data.shape[-1]
    hW = (W-1)//2

    nan_mask = np.isnan(data)
    data1 = np.where(nan_mask,0,data)    
    nan_avgs = uniff(nan_mask.astype(float),size=W, axis=-1, mode='constant')

    b_sizes = hW+1+np.arange(hW) # Boundary sizes
    count = np.hstack(( b_sizes , W*np.ones(N-2*hW), b_sizes[::-1] ))
    scale = ((count/float(W)) - nan_avgs)
    out = uniff(data1,size=W, axis=-1, mode='constant')/scale
    out = np.where(np.isclose( scale, 0), np.nan, out)
    return out

运行时测试

数据集#1:

In [807]: # Create random input array and insert NaNs
     ...: data = np.random.randint(10,size=(20,30,60)).astype(float)
     ...: 
     ...: # Add 10% NaNs across the data randomly
     ...: idx = np.random.choice(data.size,size=int(data.size*0.1),replace=0)
     ...: data.ravel()[idx] = np.nan
     ...: 
     ...: W = 5 # Window size
     ...: 

In [808]: %timeit nanmoving_mean(data,window=W,axis=2)
     ...: %timeit nanmoving_mean_numpy(data, W)
     ...: %timeit nanmoving_mean_numpy_v2(data, W)
     ...: %timeit nanmoving_mean_numpy_v3(data, W)
     ...: 
10 loops, best of 3: 22.3 ms per loop
100 loops, best of 3: 3.31 ms per loop
100 loops, best of 3: 2.99 ms per loop
1000 loops, best of 3: 1.76 ms per loop

数据集 #2 [更大的数据集]:

In [811]: # Create random input array and insert NaNs
 ...: data = np.random.randint(10,size=(120,130,160)).astype(float)
 ...: 
 ...: # Add 10% NaNs across the data randomly
 ...: idx = np.random.choice(data.size,size=int(data.size*0.1),replace=0)
 ...: data.ravel()[idx] = np.nan
 ...: 

In [812]: %timeit nanmoving_mean(data,window=W,axis=2)
     ...: %timeit nanmoving_mean_numpy(data, W)
     ...: %timeit nanmoving_mean_numpy_v2(data, W)
     ...: %timeit nanmoving_mean_numpy_v3(data, W)
     ...: 
1 loops, best of 3: 796 ms per loop
1 loops, best of 3: 486 ms per loop
1 loops, best of 3: 275 ms per loop
10 loops, best of 3: 161 ms per loop