寻找针对多维数组FFT卷积与维度乘积运算的Python高效实现方案

Nice work trying to avoid a double loop here—let's make this even more efficient and Pythonic by ditching the explicit loop entirely, leveraging NumPy's broadcasting and the linearity of convolution.

First, let's break down what your current code is doing: for each index i in the first dimension, you're tiling b[i, :, :] to match the shape of a, running the FFT convolution, summing over the first dimension, and repeating. The problem here is that tiling creates redundant copies of data, and the loop adds unnecessary Python-level overhead.

Optimized Approach 1: Remove the Loop with Broadcasting #

Instead of tiling, we can use NumPy's broadcasting to logically expand the dimensions of b and a without copying data. This lets us run fftconvolve once for all combinations, then sum the results appropriately:

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

# Expand dimensions to enable broadcasting: b -> (n, 1, m, N), a -> (1, n, m, N)
b_expanded = b[:, np.newaxis, :, :]
a_expanded = a[np.newaxis, :, :, :]

# Run FFT convolution across the last axis for all i,j pairs
convolved = fftconvolve(b_expanded, a_expanded, mode="same", axes=-1)

# Sum over the second dimension (the j index) to get (n, m, N) result
result = convolved.sum(axis=1)

# If you need the total sum over all i and j (m, N) shape:
final_result = result.sum(axis=0)

This approach:

• Eliminates the Python loop entirely, letting optimized C-level code handle all the heavy lifting
• Uses broadcasting instead of tiling, so no redundant data copies are made (saves memory)
• Is cleaner and more readable, aligning with Pythonic principles of using vectorized operations

Even Better: Use Convolution Linearity for Total Sum #

If your end goal is just the total sum over all i and j (the (m, N) result), we can use a key property of convolution: it's linear in both inputs. That means the sum of convolutions is the convolution of sums. This cuts out the n×n intermediate computation entirely:

# Sum over the first dimension for both arrays
sum_b = b.sum(axis=0)
sum_a = a.sum(axis=0)

# Convolve the summed arrays directly
final_result = fftconvolve(sum_b, sum_a, mode="same", axes=-1)

This is drastically more efficient, especially as n grows—you'll save both memory (no more n×n×m×N intermediate array) and computation time (only one convolution instead of n×n).

Verification #

To confirm these approaches match your original code (fixed to accumulate results correctly), here's the corrected original loop for comparison:

result = np.zeros((m, N), dtype=np.complex128)
for i in range(n):
    b_tiled = np.tile(b[i, :, :], (n, 1, 1)).reshape(n*m, N)
    conv = fftconvolve(b_tiled, a.reshape(n*m, N), mode="same", axes=-1).reshape(n, m, N)
    result += conv.sum(axis=0)

The final_result from either optimized approach will match this corrected loop's output exactly.

内容的提问来源于stack exchange，提问作者Zarathustra