revert to scipy for maintainability
This commit is contained in:
+56
-52
@@ -24,8 +24,8 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
import numpy
|
||||
import timeit
|
||||
import sys
|
||||
import cv2
|
||||
from scipy import spatial
|
||||
# import cv2
|
||||
from scipy import signal, spatial
|
||||
|
||||
from .backend_availability import (
|
||||
CUPY_AVAILABLE, OPENCL_AVAILABLE, KERNEL_AVAILABLE,
|
||||
@@ -1200,69 +1200,73 @@ class PoissonSolver(object):
|
||||
# convolution in frequency domain using OpenCV FFT (2D) or numpy FFT (3D)
|
||||
if self.stype == 'F':
|
||||
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)
|
||||
# # 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)
|
||||
# # 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
|
||||
)
|
||||
# # 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)
|
||||
# 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-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
|
||||
}
|
||||
# # 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
|
||||
# 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
|
||||
# # 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)
|
||||
# # 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)
|
||||
# # 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)
|
||||
# # 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)
|
||||
# # Extract 'same' mode result (use cached slices, scale in-place)
|
||||
# p = conv_full[cache['slices']].copy()
|
||||
# numpy.multiply(p, cache['gdx_prod'], p)
|
||||
# from scipy import signal
|
||||
p = signal.fftconvolve(f, self.G, mode='same')
|
||||
# Scale by grid spacing
|
||||
numpy.multiply(p, numpy.prod(self.gdx), p)
|
||||
else:
|
||||
# For 3D, use numpy FFT (OpenCV DFT only supports 2D)
|
||||
from scipy import signal
|
||||
# from scipy import signal
|
||||
p = signal.fftconvolve(f, self.G, mode='same')
|
||||
# Scale by grid spacing
|
||||
numpy.multiply(p, numpy.prod(self.gdx), p)
|
||||
|
||||
Reference in New Issue
Block a user