Блог

Я ищу эффективный способ выполнения следующей операции; вот минимальный рабочий код:

 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 ) 
 

Операция вычисляет БПФ двух массивов по типу продукта в первом индексе ( поэтому я избегаю двойного цикла по индексам диапазона(n), используя только один ).

Эталонная реализация

Я начну обертывать операцию в функцию, чтобы я мог легко сравнивать реализации

 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 в разы быстрее.