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 numpy
import timeit import timeit
import sys import sys
from scipy import signal, spatial import cv2
from scipy import spatial
from .backend_availability import ( from .backend_availability import (
CUPY_AVAILABLE, OPENCL_AVAILABLE, KERNEL_AVAILABLE, 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_crop_batch_kernel = None # Cached batch crop kernel
self._opencl_batch_mul_kernel = None # Cached batch complex multiply kernel self._opencl_batch_mul_kernel = None # Cached batch complex multiply kernel
self._opencl_last_timing = None # Last per-stage OpenCL timing in ms 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'): # if (self.ptype == 'G'):
# print( 'PoissonSolver:', self.dim, 'D grid =', self.gdims, 'with spacings', self.gdx) # print( 'PoissonSolver:', self.dim, 'D grid =', self.gdims, 'with spacings', self.gdx)
# elif (self.ptype == 'P'): # elif (self.ptype == 'P'):
@@ -1194,10 +1197,75 @@ class PoissonSolver(object):
print( self.gdims, fshape) print( self.gdims, fshape)
raise ValueError("Shape of function should match shape of grid") 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': if self.stype == 'F':
p = signal.fftconvolve(f, self.G, mode='same') if self.dim == 2:
numpy.multiply(p, numpy.prod(self.gdx), p) # 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 # convolution in spatial domain
elif self.stype == 'S': elif self.stype == 'S':