From ffb0305e108f423cbf620b225eddaf564b273613 Mon Sep 17 00:00:00 2001
From: bluepulsar <cheungken3@yahoo.com.hk>
Date: Wed, 18 Feb 2026 16:42:23 +0800
Subject: [PATCH] added opencv fft which is faster than scipy one for 2D

---
 src/orisys/HHD.py | 76 ++++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 72 insertions(+), 4 deletions(-)

diff --git a/src/orisys/HHD.py b/src/orisys/HHD.py
index 06375df..cf44110 100644
--- a/src/orisys/HHD.py
+++ b/src/orisys/HHD.py
@@ -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':