added opencv fft which is faster than scipy one for 2D

This commit is contained in:
bluepulsar
2026-02-18 16:42:23 +08:00
parent 524f750ec9
commit ffb0305e10
+72 -4
View File
@@ -24,7 +24,8 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
import numpy
import timeit
import sys
from scipy import signal, spatial
import cv2
from scipy import spatial
from .backend_availability import (
CUPY_AVAILABLE, OPENCL_AVAILABLE, KERNEL_AVAILABLE,
@@ -1013,6 +1014,8 @@ class PoissonSolver(object):
self._opencl_crop_batch_kernel = None # Cached batch crop kernel
self._opencl_batch_mul_kernel = None # Cached batch complex multiply kernel
self._opencl_last_timing = None # Last per-stage OpenCL timing in ms
# Cache for OpenCV FFT (CPU backend)
self._opencv_cache = None # Cache for OpenCV FFT buffers
# if (self.ptype == 'G'):
# print( 'PoissonSolver:', self.dim, 'D grid =', self.gdims, 'with spacings', self.gdx)
# elif (self.ptype == 'P'):
@@ -1194,10 +1197,75 @@ class PoissonSolver(object):
print( self.gdims, fshape)
raise ValueError("Shape of function should match shape of grid")
# convolution in frequency domain
# convolution in frequency domain using OpenCV FFT (2D) or numpy FFT (3D)
if self.stype == 'F':
p = signal.fftconvolve(f, self.G, mode='same')
numpy.multiply(p, numpy.prod(self.gdx), p)
if self.dim == 2:
# Use OpenCV DFT for 2D (optimized with caching)
# Convert to float32 for OpenCV (use ascontiguousarray for better performance)
f_32 = numpy.ascontiguousarray(f, dtype=numpy.float32)
# For 'same' mode convolution, pad to (2*N-1) size
full_shape = tuple(2 * s - 1 for s in f_32.shape)
# Initialize or validate cache
cache_valid = (
self._opencv_cache is not None and
self._opencv_cache['fshape'] == fshape and
self._opencv_cache['full_shape'] == full_shape
)
if not cache_valid:
# Cache G_fft (G doesn't change, so compute once)
G_32 = numpy.ascontiguousarray(self.G, dtype=numpy.float32)
G_fft = cv2.dft(G_32, flags=cv2.DFT_COMPLEX_OUTPUT)
# Pre-compute crop slices for 'same' mode extraction
# For circular convolution at size (2*H-1, 2*W-1), the center corresponds to
# position [H-1, W-1], which matches scipy's fftconvolve mode='same' extraction
starts = [fshape[i] - 1 for i in range(len(fshape))]
slices = tuple(slice(starts[i], starts[i] + fshape[i]) for i in range(len(fshape)))
# Pre-allocate reusable buffers
self._opencv_cache = {
'fshape': fshape,
'full_shape': full_shape,
'G_fft': G_fft, # Cached: G doesn't change, so compute FFT once
'slices': slices, # Cached: crop slices pre-computed
'gdx_prod': numpy.prod(self.gdx), # Cached: grid spacing product
# Reusable buffers (avoid reallocation on each call)
'f_pad': numpy.zeros(full_shape, dtype=numpy.float32), # Reused for padding
'conv_fft': numpy.empty((full_shape[0], full_shape[1], 2), dtype=numpy.float32), # Output for mulSpectrums
}
cache = self._opencv_cache
# Zero-pad f into pre-allocated buffer (faster than creating new array)
f_pad = cache['f_pad']
f_pad.fill(0.0) # Clear previous data efficiently
f_slices = tuple(slice(0, s) for s in f_32.shape)
f_pad[f_slices] = f_32
# Compute FFT of f (cv2.dft returns new array, but we can reuse for mulSpectrums)
f_fft = cv2.dft(f_pad, flags=cv2.DFT_COMPLEX_OUTPUT)
# Multiply in frequency domain using OpenCV's optimized mulSpectrums
# This is faster than manual complex multiplication and supports output buffer
# Signature: mulSpectrums(a, b, flags, c=None, conjB=False)
cv2.mulSpectrums(f_fft, cache['G_fft'], cv2.DFT_COMPLEX_OUTPUT,
cache['conv_fft'], False)
# Compute inverse FFT (cv2.idft returns new array)
conv_full = cv2.idft(cache['conv_fft'], flags=cv2.DFT_SCALE | cv2.DFT_REAL_OUTPUT)
# Extract 'same' mode result (use cached slices, scale in-place)
p = conv_full[cache['slices']].copy()
numpy.multiply(p, cache['gdx_prod'], p)
else:
# For 3D, use numpy FFT (OpenCV DFT only supports 2D)
from scipy import signal
p = signal.fftconvolve(f, self.G, mode='same')
# Scale by grid spacing
numpy.multiply(p, numpy.prod(self.gdx), p)
# convolution in spatial domain
elif self.stype == 'S':