此操作的有效pythonic方式答案

【问题标题】：Efficient pythonic way for this operation此操作的有效pythonic方式
【发布时间】：2021-06-09 12:49:35
【问题描述】：

我正在寻找一种有效的方法来执行以下操作；这是一个最小的工作代码：

import numpy as np
from scipy.signal import fftconvolve

n = 7
m = 100
N = 3000

a = np.random.rand( n,m,N ) + np.random.rand( n,m,N )*1j
b = np.random.rand( n,m,N ) + np.random.rand( n,m,N )*1j

# we want product over the n-dimension, with fftconvolve in the m-indices elementwise

old_shape= a.shape 
new_shape= ( n*m, N )
    
a  = a.reshape( new_shape ) 

for i in range( n ):
    
    b_tiled = np.tile( b[ i, :, : ], ( n, 1, 1 )).reshape( new_shape )
    result  = ( fftconvolve( b_tiled, a, mode="same", axes=-1 ) ).reshape( old_shape )
    
    result  = result.sum( axis=0 )

该操作在第一个索引中以类似乘积的方式计算两个数组的 FFT（因此我避免了 range(n) 索引上的双循环，仅使用一个）。

【问题讨论】：

标签： python-3.x numpy scipy

【解决方案1】：

参考实现

我将开始将操作包装在一个函数中，以便我可以轻松比较实现

def ref_impl(a,b):
    n,m,N = a.shape
    a  = a.reshape( m*n, N ) 
    ans = []
    for i in range( n ):
        b_tiled = np.tile( b[ i, :, : ], ( n, 1, 1 )).reshape( n*m, N )
        result  = ( fftconvolve( b_tiled, a, mode="same", axes=-1 ) ).reshape( n, m, N )
        ans.append(result.sum( axis=0 ))
    return np.array(ans);

简化平铺数组的总和

平铺复制一个元素，总和减少一个轴上的所有元素。

注意b_tiled 在第一个轴上是恒定的，你有一些整形但一切都对齐了，所以第一个result[k] = fftconvolve(b[i,:,:], a[k, :, :]) 所以第二个结果可以计算为

result = sum(fftconvolve(b[i,:,:], a[k, :, :]) for k in range(n))

因为 fftconvolve 是线性的，所以可以写成

result = fftconvolve(b[i,:,:], sum(a[k, :, :]for k in range(n)))

而且这种形式可以矢量化

def impl1(a,b):
    n,m,N = a.shape
    ans = []
    for i in range( n ):
        result  = ( fftconvolve( b[i, :, :], a.sum(axis=0), mode="same", axes=-1 ) )
        ans.append(result)
    return np.array(ans);

简化切片堆栈

索引在一个轴上取一个元素，栈会构造一个包含多个元素的数组，这个操作可以用单个向量化操作代替

def impl2(a,b):
    n,m,N = a.shape
    return ( fftconvolve( b, np.tile(a.sum(axis=0), (n, 1, 1)), mode="same", axes=-1 ) )

复杂性

提议的实现将使用更少的暂存空间，将摆脱 for 循环，并且将只执行一个 fftconvolve 而不是 n。总之，对于大型输入，它将运行 n 倍快。

【讨论】：