refractor HHD
This commit is contained in:
ykcheungab
-2072
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,20 @@
|
||||
"""
|
||||
HHD subsystem package.
|
||||
|
||||
This package is intended to hold the refactored implementation currently living in
|
||||
`orisys.HHD`. The old monolithic file can remain as a compatibility layer while code
|
||||
is moved here gradually.
|
||||
|
||||
Suggested migration target:
|
||||
- `api.py` -> public `nHHD` class
|
||||
- `grid.py` -> `StructuredGrid`
|
||||
- `poisson.py` -> `PoissonSolver`
|
||||
- `timing.py` -> `Timer`
|
||||
- `kernels.py` -> helper kernels / backend-specific math
|
||||
"""
|
||||
from .api import nHHD
|
||||
from .grid import StructuredGrid
|
||||
from .poisson import PoissonSolver
|
||||
from .timing import Timer
|
||||
|
||||
__all__ = ["nHHD", "StructuredGrid", "PoissonSolver", "Timer"]
|
||||
@@ -0,0 +1,456 @@
|
||||
import numpy
|
||||
|
||||
from ..backend_availability import (
|
||||
CUPY_AVAILABLE,
|
||||
OPENCL_AVAILABLE,
|
||||
KERNEL_AVAILABLE,
|
||||
cp,
|
||||
_torque_kernel,
|
||||
)
|
||||
from .grid import StructuredGrid
|
||||
from .poisson import PoissonSolver
|
||||
|
||||
|
||||
class nHHD(object):
|
||||
|
||||
# constructor
|
||||
def __init__(self, cuda=True, **kwargs):
|
||||
# Select Poisson backend: 'auto' (default), 'cupy', 'opencl', or 'cpu'
|
||||
backend = kwargs.get('poisson_backend', 'auto').lower()
|
||||
if backend not in ('auto', 'cupy', 'opencl', 'cpu'):
|
||||
raise ValueError("poisson_backend must be one of: 'auto', 'cupy', 'opencl', 'cpu'")
|
||||
|
||||
# Resolve 'auto' to a specific backend
|
||||
if backend == 'auto':
|
||||
backend = 'cupy' if CUPY_AVAILABLE else ('opencl' if OPENCL_AVAILABLE else 'cpu')
|
||||
|
||||
# Validate backend availability and fallback to 'cpu' if needed
|
||||
if backend == 'cupy' and not CUPY_AVAILABLE:
|
||||
backend = 'cpu'
|
||||
elif backend == 'opencl' and not OPENCL_AVAILABLE:
|
||||
backend = 'cpu'
|
||||
|
||||
self.debug = kwargs.get('debug', False)
|
||||
self.poisson_backend = backend
|
||||
|
||||
# Initialize CUDA streams for cupy backend
|
||||
if self.poisson_backend == 'cupy':
|
||||
self.stream1 = cp.cuda.Stream()
|
||||
self.stream2 = cp.cuda.Stream()
|
||||
|
||||
self.cuda = cuda
|
||||
|
||||
args = kwargs.keys()
|
||||
|
||||
if ('grid' in args) == ('points' in args):
|
||||
raise SyntaxError("nHHD Solver needs either shape of a regular grid or the points in an rectilinear/unstructured grid")
|
||||
|
||||
if ('grid' in args) != ('spacings' in args):
|
||||
raise SyntaxError("nHHD needs spacings for regular grid")
|
||||
|
||||
if 'grid' in args:
|
||||
self.mesh = StructuredGrid(grid=kwargs['grid'], spacings=kwargs['spacings'])
|
||||
self.psolver = PoissonSolver(grid=self.mesh.dims, spacings=self.mesh.dx, poisson_backend=self.poisson_backend)
|
||||
|
||||
self.dim = self.mesh.dim
|
||||
if (self.dim != 2) and (self.dim != 3):
|
||||
raise ValueError("nHHD Solver works only for 2 and 3 dimensions")
|
||||
|
||||
self._force_cpu_cache = None
|
||||
self.psolver.prepare(False)
|
||||
|
||||
def get_force(self, vfield, verbose=False):
|
||||
if self.poisson_backend == 'cupy':
|
||||
return self.get_force_cuda(vfield, verbose)
|
||||
else:
|
||||
return self.get_force_cpu(vfield, verbose)
|
||||
|
||||
def get_shear(self):
|
||||
if self.poisson_backend == 'cpu' or self.poisson_backend == 'opencl':
|
||||
return self.h
|
||||
elif self.poisson_backend == 'cupy':
|
||||
# Use CPU copy from get_force_cuda when set (stable snapshot; avoids jumps)
|
||||
if hasattr(self, 'h') and self.h is not None:
|
||||
return self.h
|
||||
return cp.asnumpy(self._h_gpu)
|
||||
|
||||
def get_normal(self):
|
||||
if self.poisson_backend == 'cpu' or self.poisson_backend == 'opencl':
|
||||
return self.d
|
||||
elif self.poisson_backend == 'cupy':
|
||||
# Use CPU copy from get_force_cuda when set (stable snapshot; avoids jumps)
|
||||
if hasattr(self, 'd') and self.d is not None:
|
||||
return self.d
|
||||
return cp.asnumpy(self._d_gpu)
|
||||
|
||||
def _solve_poisson_backend(self, field, verbose=False):
|
||||
"""
|
||||
Dispatch Poisson solve using the selected backend.
|
||||
"""
|
||||
if self.poisson_backend == 'opencl':
|
||||
return self.psolver.solve_opencl(field, verbose, profile=self.debug)
|
||||
|
||||
if self.poisson_backend == 'cupy':
|
||||
return self.psolver.solve_gpu(field, verbose)
|
||||
|
||||
if self.poisson_backend == 'cpu':
|
||||
return self.psolver.solve(field, verbose)
|
||||
|
||||
# create the 3-component decomposition
|
||||
def decompose(self, vfield, verbose=0):
|
||||
# decompose for cpu and opencl backend
|
||||
|
||||
if vfield.shape[-1] != self.dim:
|
||||
raise ValueError("nHHD.decompose requires a valid-dimensional vector field")
|
||||
|
||||
if verbose > 0:
|
||||
print('\n ==> nhhd.decompose', vfield.shape)
|
||||
|
||||
if self.dim == 2:
|
||||
|
||||
# compute div curl
|
||||
(self.div, self.curlw) = self.mesh.divcurl(vfield)
|
||||
|
||||
if self.poisson_backend == 'opencl':
|
||||
self.div = numpy.ascontiguousarray(self.div, dtype=numpy.float32)
|
||||
self.curlw = numpy.ascontiguousarray(self.curlw, dtype=numpy.float32)
|
||||
|
||||
if verbose > 0:
|
||||
print('div =', self.div.shape, self.div.min(), self.div.max())
|
||||
print('curlw =', self.curlw.shape, self.curlw.min(), self.curlw.max())
|
||||
|
||||
# compute potentials
|
||||
if self.poisson_backend == 'opencl' and ('OPENCL_AVAILABLE' in globals() and OPENCL_AVAILABLE):
|
||||
self.nD, self.nRu = self.psolver.solve_opencl_batch2(self.div, self.curlw, verbose > 0, profile=self.debug)
|
||||
else:
|
||||
self.nD = self._solve_poisson_backend(self.div, verbose > 0)
|
||||
self.nRu = self._solve_poisson_backend(self.curlw, verbose > 0)
|
||||
if verbose > 0:
|
||||
print('nD =', self.nD.shape, self.nD.min(), self.nD.max())
|
||||
print('nR =', self.nRu.shape, self.nRu.min(), self.nRu.max())
|
||||
|
||||
# compute fields as gradients of potentials
|
||||
self.d = self.mesh.gradient(self.nD)
|
||||
if verbose > 0:
|
||||
print('d =', self.d.shape, self.d.min(), self.d.max())
|
||||
|
||||
self.r = self.mesh.rotated_gradient(self.nRu)
|
||||
if verbose > 0:
|
||||
print('r =', self.r.shape, self.r.min(), self.r.max())
|
||||
else:
|
||||
|
||||
# compute div curl
|
||||
(self.div, self.curlu, self.curlv, self.curlw) = self.mesh.divcurl(vfield)
|
||||
|
||||
if verbose > 0:
|
||||
print('div =', self.div.shape, self.div.min(), self.div.max())
|
||||
print('curlu =', self.curlu.shape, self.curlu.min(), self.curlu.max())
|
||||
print('curlv =', self.curlv.shape, self.curlv.min(), self.curlv.max())
|
||||
print('curlw =', self.curlw.shape, self.curlw.min(), self.curlw.max())
|
||||
|
||||
# compute potentials
|
||||
self.nD = self._solve_poisson_backend(self.div, verbose > 0)
|
||||
if verbose > 0:
|
||||
print('nD =', self.nD.shape, self.nD.min(), self.nD.max())
|
||||
|
||||
self.nRu = self._solve_poisson_backend(self.curlu, verbose > 0)
|
||||
self.nRu *= -1.0
|
||||
if verbose > 0:
|
||||
print('nRu =', self.nRu.shape, self.nRu.min(), self.nRu.max())
|
||||
|
||||
self.nRv = self._solve_poisson_backend(self.curlv, verbose > 0)
|
||||
self.nRv *= -1.0
|
||||
if verbose > 0:
|
||||
print('nRv =', self.nRv.shape, self.nRv.min(), self.nRv.max())
|
||||
|
||||
self.nRw = self._solve_poisson_backend(self.curlw, verbose > 0)
|
||||
self.nRw *= -1.0
|
||||
|
||||
if verbose > 0:
|
||||
print('nRw =', self.nRw.shape, self.nRw.min(), self.nRw.max())
|
||||
|
||||
# compute divergent field as gradient of D potential
|
||||
self.d = self.mesh.gradient(self.nD)
|
||||
if verbose > 0:
|
||||
print('d =', self.d.shape, self.d.min(), self.d.max())
|
||||
|
||||
# temporary use of r
|
||||
self.r = numpy.stack((self.nRu, self.nRv, self.nRw), axis=-1)
|
||||
|
||||
# compute rotational field as curl of R potential
|
||||
(cu, cv, cw) = self.mesh.curl3D(self.r)
|
||||
self.r = numpy.stack((cu, cv, cw), axis=-1)
|
||||
if verbose > 0:
|
||||
print('r =', self.r.shape, self.r.min(), self.r.max())
|
||||
|
||||
self.h = numpy.add(self.d, self.r)
|
||||
numpy.subtract(vfield, self.h, self.h)
|
||||
|
||||
if verbose > 0:
|
||||
print('h =', self.h.shape, self.h.min(), self.h.max())
|
||||
|
||||
def decompose_cuda(self, vfield, verbose=0):
|
||||
"""
|
||||
GPU-optimized decomposition that keeps intermediate results on GPU.
|
||||
Returns GPU arrays to avoid CPU-GPU transfers.
|
||||
|
||||
Args:
|
||||
vfield: Vector field (numpy array or CuPy array)
|
||||
verbose: Verbosity level (default: 0)
|
||||
|
||||
Returns:
|
||||
Tuple of GPU arrays: (d_gpu, h_gpu, r_gpu, nRu_gpu) for 2D fields
|
||||
"""
|
||||
if not CUPY_AVAILABLE:
|
||||
if verbose > 0:
|
||||
print("CuPy not available, falling back to CPU decompose")
|
||||
self.decompose(vfield, verbose)
|
||||
return None
|
||||
|
||||
if vfield.shape[-1] != self.dim:
|
||||
raise ValueError("nHHD.decompose requires a valid-dimensional vector field")
|
||||
|
||||
if verbose > 0:
|
||||
print('\n ==> nhhd.decompose_cuda', vfield.shape)
|
||||
|
||||
# Convert input to GPU if needed
|
||||
if isinstance(vfield, numpy.ndarray):
|
||||
vfield_gpu = cp.asarray(vfield, dtype=cp.float32)
|
||||
else:
|
||||
vfield_gpu = vfield
|
||||
|
||||
if self.dim == 2:
|
||||
# Compute div curl on GPU
|
||||
dx0, dx1 = self.mesh.dx[0], self.mesh.dx[1]
|
||||
|
||||
dudy, dudx = cp.gradient(vfield_gpu[:, :, 0], dx0, dx1)
|
||||
dvdy, dvdx = cp.gradient(vfield_gpu[:, :, 1], dx0, dx1)
|
||||
|
||||
div_gpu = cp.add(dudx, dvdy)
|
||||
curlw_gpu = cp.subtract(dvdx, dudy)
|
||||
|
||||
if verbose > 0:
|
||||
print('div =', div_gpu.shape, float(div_gpu.min()), float(div_gpu.max()))
|
||||
print('curlw =', curlw_gpu.shape, float(curlw_gpu.min()), float(curlw_gpu.max()))
|
||||
|
||||
# Use a single stream for both solves. Running concurrent FFT-based
|
||||
# solves on separate streams can cause occasional long-run jitter.
|
||||
nD_gpu = self.psolver._solve_gpu_internal(div_gpu)
|
||||
nRu_gpu = self.psolver._solve_gpu_internal(curlw_gpu)
|
||||
cp.cuda.get_current_stream().synchronize()
|
||||
|
||||
if verbose > 0:
|
||||
print('nD =', nD_gpu.shape, float(nD_gpu.min()), float(nD_gpu.max()))
|
||||
print('nR =', nRu_gpu.shape, float(nRu_gpu.min()), float(nRu_gpu.max()))
|
||||
|
||||
# Compute gradients on GPU
|
||||
ddy, ddx = cp.gradient(nD_gpu, dx0, dx1)
|
||||
d_gpu = cp.stack((ddx, ddy), axis=-1)
|
||||
|
||||
if verbose > 0:
|
||||
print('d =', d_gpu.shape, float(d_gpu.min()), float(d_gpu.max()))
|
||||
|
||||
# Compute rotated gradient on GPU
|
||||
rdy, rdx = cp.gradient(nRu_gpu, dx0, dx1)
|
||||
rdy = cp.multiply(rdy, -1.0)
|
||||
r_gpu = cp.stack((rdy, rdx), axis=-1)
|
||||
|
||||
if verbose > 0:
|
||||
print('r =', r_gpu.shape, float(r_gpu.min()), float(r_gpu.max()))
|
||||
|
||||
# Compute h on GPU
|
||||
h_gpu = cp.add(d_gpu, r_gpu)
|
||||
h_gpu = cp.subtract(vfield_gpu, h_gpu)
|
||||
|
||||
if verbose > 0:
|
||||
print('h =', h_gpu.shape, float(h_gpu.min()), float(h_gpu.max()))
|
||||
|
||||
# Ensure all default-stream work (gradient, etc.) is complete before use
|
||||
cp.cuda.get_current_stream().synchronize()
|
||||
|
||||
# Store GPU arrays
|
||||
self._d_gpu = d_gpu
|
||||
self._h_gpu = h_gpu
|
||||
self._r_gpu = r_gpu
|
||||
self._nRu_gpu = nRu_gpu
|
||||
|
||||
return (d_gpu, h_gpu, r_gpu, nRu_gpu)
|
||||
|
||||
else:
|
||||
# For 3D, fall back to CPU for now
|
||||
if verbose > 0:
|
||||
print("3D GPU decompose not yet implemented, using CPU")
|
||||
self.decompose(vfield, verbose)
|
||||
return None
|
||||
|
||||
def get_force_cpu(self, vfield, verbose=False):
|
||||
shape = (vfield.shape[0], vfield.shape[1])
|
||||
if (not hasattr(self, '_force_cpu_cache') or
|
||||
self._force_cpu_cache is None or
|
||||
self._force_cpu_cache['shape'] != shape):
|
||||
self._init_force_cpu_cache(shape)
|
||||
grid = self._force_cpu_cache['grid']
|
||||
|
||||
self.decompose(vfield, verbose)
|
||||
|
||||
# Optimized: use abs instead of absolute (slightly faster)
|
||||
self.Sn = numpy.sum(numpy.abs(self.d))
|
||||
|
||||
# Optimized: direct indexing and sum (already efficient)
|
||||
self.Stx = numpy.sum(self.h[:, :, 0])
|
||||
self.Sty = numpy.sum(self.h[:, :, 1])
|
||||
|
||||
# Optimized: find both max and min indices in single pass through flattened array
|
||||
nRu_flat = self.nRu.ravel()
|
||||
nR_max_idx = numpy.argmax(nRu_flat)
|
||||
nR_min_idx = numpy.argmin(nRu_flat)
|
||||
# Convert flat indices to 2D coordinates (faster than unravel_index)
|
||||
nR_max_loc = (nR_max_idx // shape[1], nR_max_idx % shape[1])
|
||||
nR_min_loc = (nR_min_idx // shape[1], nR_min_idx % shape[1])
|
||||
|
||||
# Optimized: compute cross products directly without creating intermediate arrays
|
||||
# For 2D: cross(a, b) = a[0]*b[1] - a[1]*b[0]
|
||||
# Extract components once for reuse
|
||||
r_x = self.r[:, :, 0]
|
||||
r_y = self.r[:, :, 1]
|
||||
grid_x = grid[:, :, 0]
|
||||
grid_y = grid[:, :, 1]
|
||||
|
||||
# Compute cross products directly (faster than numpy.cross for 2D)
|
||||
# cross(r_p, r) where r_p = grid - nR_max_loc
|
||||
max_y, max_x = nR_max_loc[0], nR_max_loc[1]
|
||||
cross_p = (grid_x - max_x) * r_y - (grid_y - max_y) * r_x
|
||||
|
||||
# cross(r_n, r) where r_n = grid - nR_min_loc
|
||||
min_y, min_x = nR_min_loc[0], nR_min_loc[1]
|
||||
cross_n = (grid_x - min_x) * r_y - (grid_y - min_y) * r_x
|
||||
|
||||
self.Stor = numpy.sum(cross_p) + numpy.sum(cross_n)
|
||||
|
||||
# Optimized: pre-allocate output array with explicit dtype
|
||||
output = numpy.array([self.Sn, self.Stx, self.Sty, self.Stor], dtype=numpy.float32)
|
||||
|
||||
# Optimized: combine threshold checks (short-circuit evaluation)
|
||||
flag = (self.Sn >= 300 and self.Stx >= 300 and self.Sty >= 300 and numpy.abs(self.Stor) >= 1000)
|
||||
|
||||
return flag, output
|
||||
|
||||
def _init_force_cpu_cache(self, shape):
|
||||
"""
|
||||
Initialize cached CPU grid arrays for force computation.
|
||||
Avoids per-frame mgrid/stack allocations and reduces jitter.
|
||||
"""
|
||||
mgrid = numpy.mgrid[0:shape[0], 0:shape[1]]
|
||||
self._force_cpu_cache = {
|
||||
'shape': shape,
|
||||
'grid': numpy.stack((mgrid[0], mgrid[1]), axis=2),
|
||||
}
|
||||
|
||||
def _init_force_gpu_cache(self, shape):
|
||||
"""
|
||||
Initialize cached GPU arrays for force computation.
|
||||
Call once when shape is known to avoid repeated allocations.
|
||||
|
||||
Args:
|
||||
shape: Tuple (height, width) of the vector field
|
||||
"""
|
||||
if not CUPY_AVAILABLE:
|
||||
return
|
||||
|
||||
# Cache grid arrays (only need to create once per shape)
|
||||
self._force_gpu_cache = {
|
||||
'shape': shape,
|
||||
'mgrid_0': cp.arange(shape[0], dtype=cp.float32)[:, None],
|
||||
'mgrid_1': cp.arange(shape[1], dtype=cp.float32)[None, :],
|
||||
}
|
||||
# Pre-broadcast for efficiency
|
||||
self._force_gpu_cache['mgrid_0_bc'] = cp.broadcast_to(
|
||||
self._force_gpu_cache['mgrid_0'], shape
|
||||
).copy() # .copy() to make contiguous
|
||||
self._force_gpu_cache['mgrid_1_bc'] = cp.broadcast_to(
|
||||
self._force_gpu_cache['mgrid_1'], shape
|
||||
).copy()
|
||||
|
||||
def get_force_cuda(self, vfield, verbose=False):
|
||||
"""
|
||||
Further optimized GPU force computation using CuPy ElementwiseKernel.
|
||||
Fuses all operations into a single kernel for maximum performance.
|
||||
|
||||
Args:
|
||||
vfield: Vector field (numpy array)
|
||||
verbose: Print debug information (default: False)
|
||||
|
||||
Returns:
|
||||
tuple: (flag, output) where:
|
||||
- flag: Boolean indicating if force thresholds are met
|
||||
- output: Array [Sn, Stx, Sty, Stor] containing force components
|
||||
"""
|
||||
if not CUPY_AVAILABLE:
|
||||
if verbose:
|
||||
print("CuPy not available, falling back to CPU get_force")
|
||||
return self.get_force(vfield, verbose)
|
||||
|
||||
# Decompose the vector field
|
||||
self.decompose_cuda(vfield, verbose)
|
||||
|
||||
d_gpu = self._d_gpu
|
||||
h_gpu = self._h_gpu
|
||||
r_gpu = self._r_gpu
|
||||
nRu_gpu = self._nRu_gpu
|
||||
|
||||
# Use reduction kernels for better performance
|
||||
# Compute Sn as L1 sum directly on GPU (faster than linalg.norm for this case)
|
||||
Sn = float(cp.sum(cp.abs(d_gpu)))
|
||||
|
||||
# Compute Stx, Sty
|
||||
Stx = float(cp.sum(h_gpu[:, :, 0]))
|
||||
Sty = float(cp.sum(h_gpu[:, :, 1]))
|
||||
|
||||
# Find extrema indices
|
||||
nR_max_idx = int(cp.argmax(nRu_gpu))
|
||||
nR_min_idx = int(cp.argmin(nRu_gpu))
|
||||
|
||||
shape = nRu_gpu.shape
|
||||
max_0, max_1 = nR_max_idx // shape[1], nR_max_idx % shape[1]
|
||||
min_0, min_1 = nR_min_idx // shape[1], nR_min_idx % shape[1]
|
||||
|
||||
# Initialize cache if needed
|
||||
if not hasattr(self, '_force_gpu_cache') or self._force_gpu_cache is None or self._force_gpu_cache['shape'] != shape:
|
||||
self._init_force_gpu_cache(shape)
|
||||
|
||||
mgrid_0 = self._force_gpu_cache['mgrid_0_bc']
|
||||
mgrid_1 = self._force_gpu_cache['mgrid_1_bc']
|
||||
|
||||
r_x = r_gpu[:, :, 0]
|
||||
r_y = r_gpu[:, :, 1]
|
||||
|
||||
if KERNEL_AVAILABLE and _torque_kernel is not None:
|
||||
torque = _torque_kernel(
|
||||
r_x, r_y, mgrid_0, mgrid_1,
|
||||
float(max_0), float(max_1), float(min_0), float(min_1)
|
||||
)
|
||||
Stor = float(cp.sum(torque))
|
||||
else:
|
||||
Stor = float(cp.sum(
|
||||
(mgrid_0 - max_0) * r_y - (mgrid_1 - max_1) * r_x +
|
||||
(mgrid_0 - min_0) * r_y - (mgrid_1 - min_1) * r_x
|
||||
))
|
||||
|
||||
self.Sn, self.Stx, self.Sty, self.Stor = Sn, Stx, Sty, Stor
|
||||
output = numpy.array([Sn, Stx, Sty, Stor], dtype=numpy.float32)
|
||||
|
||||
# Keep stable CPU copies for get_normal()/get_shear() so callers get this frame's
|
||||
# result instead of reading GPU memory that can change (avoids CuPy-only jumps)
|
||||
self.d = cp.asnumpy(self._d_gpu)
|
||||
self.h = cp.asnumpy(self._h_gpu)
|
||||
|
||||
flag = (Sn >= 300 and Stx >= 300 and Sty >= 300 and abs(Stor) >= 1000)
|
||||
|
||||
return flag, output
|
||||
|
||||
def clear_force_gpu_cache(self):
|
||||
"""Clear the GPU cache for force computation to free memory."""
|
||||
if hasattr(self, '_force_gpu_cache'):
|
||||
self._force_gpu_cache = None
|
||||
if CUPY_AVAILABLE:
|
||||
cp.get_default_memory_pool().free_all_blocks()
|
||||
@@ -0,0 +1,385 @@
|
||||
import numpy
|
||||
import timeit
|
||||
import sys
|
||||
# import cv2
|
||||
from scipy import signal, spatial
|
||||
|
||||
from ..backend_availability import (
|
||||
CUPY_AVAILABLE, OPENCL_AVAILABLE, KERNEL_AVAILABLE,
|
||||
cp, cp_signal, _torque_kernel,
|
||||
cl, cla, cl_rfftn, cl_irfftn
|
||||
)
|
||||
from .timing import Timer
|
||||
|
||||
class StructuredGrid(object):
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# --------------------------------------------------------------------------
|
||||
def __init__(self, **kwargs):
|
||||
|
||||
'''
|
||||
kwargs:
|
||||
grid: ndarray of grid dimensions (Y,X) or (Z,Y,X)
|
||||
spacings: ndarray of grid spacings (dy, dx) or (dz, dy, dx)
|
||||
verbose: verbosity level
|
||||
'''
|
||||
|
||||
args = kwargs.keys()
|
||||
|
||||
if ('grid' not in args) or ('spacings' not in args):
|
||||
raise SyntaxError("Dimensions and spacings of the grid are required")
|
||||
|
||||
verbose = 0
|
||||
if 'verbose' in args:
|
||||
verbose = kwargs['verbose']
|
||||
|
||||
self.dims = kwargs['grid']
|
||||
self.dx = kwargs['spacings']
|
||||
self.dim = len(self.dims)
|
||||
|
||||
if self.dim != 2 and self.dim != 3:
|
||||
raise ValueError("StructuredGrid works for 2D and 3D only")
|
||||
|
||||
if self.dim != len(self.dx):
|
||||
raise ValueError("Dimensions of spacings should match that of the grid")
|
||||
|
||||
# Initialize GPU cache
|
||||
self._gpu_divcurl_cache = None
|
||||
|
||||
if verbose > 0:
|
||||
print(' Initializing', self.dim, 'D structured grid...', end ='')
|
||||
sys.stdout.flush()
|
||||
mtimer = Timer()
|
||||
print(' Done!', end ='')
|
||||
mtimer.end()
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# --------------------------------------------------------------------------
|
||||
def divcurl(self, vfield, verbose=False):
|
||||
# compute div and curl of vfield (vector field)
|
||||
#if (vfield.shape[-1] != self.dim) or (vfield.shape[0:self.dim] - self.dims).any():
|
||||
if (vfield.shape[-1] != self.dim) or (vfield.shape[0:self.dim] != self.dims):
|
||||
raise ValueError("Dimensions of vector field should match that of the grid")
|
||||
|
||||
if verbose:
|
||||
print(' Computing div and curl...', end ='')
|
||||
sys.stdout.flush()
|
||||
mtimer = Timer()
|
||||
|
||||
if self.dim == 2:
|
||||
|
||||
# self.dx = (dy,dx)
|
||||
dudy, dudx = numpy.gradient(vfield[:,:,0], self.dx[0], self.dx[1])
|
||||
dvdy, dvdx = numpy.gradient(vfield[:,:,1], self.dx[0], self.dx[1])
|
||||
|
||||
numpy.add(dudx, dvdy, dudx)
|
||||
numpy.subtract(dvdx, dudy, dvdx)
|
||||
|
||||
if verbose:
|
||||
print(' Done!', end ='')
|
||||
mtimer.end()
|
||||
|
||||
return (dudx, dvdx)
|
||||
|
||||
elif self.dim == 3:
|
||||
|
||||
# self.dx = (dz,dy,dx)
|
||||
dudz, dudy, dudx = numpy.gradient(vfield[:,:,:,0], self.dx[0], self.dx[1], self.dx[2])
|
||||
dvdz, dvdy, dvdx = numpy.gradient(vfield[:,:,:,1], self.dx[0], self.dx[1], self.dx[2])
|
||||
dwdz, dwdy, dwdx = numpy.gradient(vfield[:,:,:,2], self.dx[0], self.dx[1], self.dx[2])
|
||||
|
||||
numpy.add(dudx, dvdy, dudx)
|
||||
numpy.add(dudx, dvdz, dudx)
|
||||
|
||||
numpy.subtract(dwdy, dvdz, dwdy)
|
||||
numpy.subtract(dudz, dwdx, dudz)
|
||||
numpy.subtract(dvdx, dudy, dvdx)
|
||||
|
||||
if verbose:
|
||||
print(' Done!', end ='')
|
||||
mtimer.end()
|
||||
|
||||
return (dudx, dwdy, dudz, dvdx)
|
||||
|
||||
# def divcurl_gpu(self, vfield, verbose=False):
|
||||
# """
|
||||
# GPU-accelerated computation of divergence and curl of vector field.
|
||||
# Uses CuPy for faster gradient computation on GPU.
|
||||
# """
|
||||
# if not CUPY_AVAILABLE:
|
||||
# if verbose:
|
||||
# print('CuPy not available, falling back to CPU')
|
||||
# return self.divcurl(vfield, verbose)
|
||||
|
||||
# if (vfield.shape[-1] != self.dim) or (vfield.shape[0:self.dim] != self.dims):
|
||||
# raise ValueError("Dimensions of vector field should match that of the grid")
|
||||
|
||||
# if verbose:
|
||||
# print(' Computing div and curl on GPU...', end ='')
|
||||
# sys.stdout.flush()
|
||||
# mtimer = Timer()
|
||||
|
||||
# if self.dim == 2:
|
||||
# # Transfer to GPU
|
||||
# vfield_gpu = cp.asarray(vfield, dtype=cp.float32)
|
||||
|
||||
# # Compute gradients on GPU
|
||||
# dudy, dudx = cp.gradient(vfield_gpu[:,:,0], self.dx[0], self.dx[1])
|
||||
# dvdy, dvdx = cp.gradient(vfield_gpu[:,:,1], self.dx[0], self.dx[1])
|
||||
|
||||
# # Compute div and curl on GPU
|
||||
# div_gpu = dudx + dvdy
|
||||
# curl_gpu = dvdx - dudy
|
||||
|
||||
# # Transfer back to CPU
|
||||
# div = cp.asnumpy(div_gpu)
|
||||
# curl = cp.asnumpy(curl_gpu)
|
||||
|
||||
# if verbose:
|
||||
# print(' Done!', end ='')
|
||||
# mtimer.end()
|
||||
|
||||
# return (div, curl)
|
||||
|
||||
# elif self.dim == 3:
|
||||
# # Transfer to GPU
|
||||
# vfield_gpu = cp.asarray(vfield, dtype=cp.float32)
|
||||
|
||||
# # Compute gradients on GPU
|
||||
# dudz, dudy, dudx = cp.gradient(vfield_gpu[:,:,:,0], self.dx[0], self.dx[1], self.dx[2])
|
||||
# dvdz, dvdy, dvdx = cp.gradient(vfield_gpu[:,:,:,1], self.dx[0], self.dx[1], self.dx[2])
|
||||
# dwdz, dwdy, dwdx = cp.gradient(vfield_gpu[:,:,:,2], self.dx[0], self.dx[1], self.dx[2])
|
||||
|
||||
# # Compute div and curl components on GPU
|
||||
# div_gpu = dudx + dvdy + dwdz
|
||||
# curlu_gpu = dwdy - dvdz
|
||||
# curlv_gpu = dudz - dwdx
|
||||
# curlw_gpu = dvdx - dudy
|
||||
|
||||
# # Transfer back to CPU
|
||||
# div = cp.asnumpy(div_gpu)
|
||||
# curlu = cp.asnumpy(curlu_gpu)
|
||||
# curlv = cp.asnumpy(curlv_gpu)
|
||||
# curlw = cp.asnumpy(curlw_gpu)
|
||||
|
||||
# if verbose:
|
||||
# print(' Done!', end ='')
|
||||
# mtimer.end()
|
||||
|
||||
# return (div, curlu, curlv, curlw)
|
||||
|
||||
def curl3D(self, vfield, verbose=False):
|
||||
|
||||
#if (vfield.shape[-1] != self.dim) or (vfield.shape[0:self.dim] - self.dims).any():
|
||||
if (vfield.shape[-1] != self.dim) or (vfield.shape[0:self.dim] != self.dims):
|
||||
raise ValueError("Dimensions of vector field should match that of the grid")
|
||||
|
||||
if self.dim != 3:
|
||||
raise ValueError("curl3D works only for 2D")
|
||||
|
||||
if verbose:
|
||||
print(' Computing curl...', end ='')
|
||||
sys.stdout.flush()
|
||||
mtimer = Timer()
|
||||
|
||||
# self.dx = (dz,dy,dx)
|
||||
dudz, dudy, dudx = numpy.gradient(vfield[:,:,:,0], self.dx[0], self.dx[1], self.dx[2])
|
||||
dvdz, dvdy, dvdx = numpy.gradient(vfield[:,:,:,1], self.dx[0], self.dx[1], self.dx[2])
|
||||
dwdz, dwdy, dwdx = numpy.gradient(vfield[:,:,:,2], self.dx[0], self.dx[1], self.dx[2])
|
||||
|
||||
numpy.subtract(dwdy, dvdz, dwdy)
|
||||
numpy.subtract(dudz, dwdx, dudz)
|
||||
numpy.subtract(dvdx, dudy, dvdx)
|
||||
|
||||
if verbose:
|
||||
print(' Done!', end ='')
|
||||
mtimer.end()
|
||||
|
||||
return (dwdy, dudz, dvdx)
|
||||
|
||||
def curl3D_gpu(self, vfield, verbose=False):
|
||||
"""
|
||||
GPU-accelerated computation of curl for 3D vector field.
|
||||
Uses CuPy for faster gradient computation on GPU.
|
||||
"""
|
||||
if not CUPY_AVAILABLE:
|
||||
if verbose:
|
||||
print('CuPy not available, falling back to CPU')
|
||||
return self.curl3D(vfield, verbose)
|
||||
|
||||
if (vfield.shape[-1] != self.dim) or (vfield.shape[0:self.dim] != self.dims):
|
||||
raise ValueError("Dimensions of vector field should match that of the grid")
|
||||
|
||||
if self.dim != 3:
|
||||
raise ValueError("curl3D works only for 3D")
|
||||
|
||||
if verbose:
|
||||
print(' Computing curl on GPU...', end ='')
|
||||
sys.stdout.flush()
|
||||
mtimer = Timer()
|
||||
|
||||
# Transfer to GPU
|
||||
vfield_gpu = cp.asarray(vfield, dtype=cp.float32)
|
||||
|
||||
# Compute gradients on GPU
|
||||
dudz, dudy, dudx = cp.gradient(vfield_gpu[:,:,:,0], self.dx[0], self.dx[1], self.dx[2])
|
||||
dvdz, dvdy, dvdx = cp.gradient(vfield_gpu[:,:,:,1], self.dx[0], self.dx[1], self.dx[2])
|
||||
dwdz, dwdy, dwdx = cp.gradient(vfield_gpu[:,:,:,2], self.dx[0], self.dx[1], self.dx[2])
|
||||
|
||||
# Compute curl components on GPU
|
||||
curl_u = dwdy - dvdz
|
||||
curl_v = dudz - dwdx
|
||||
curl_w = dvdx - dudy
|
||||
|
||||
# Transfer back to CPU
|
||||
curlu = cp.asnumpy(curl_u)
|
||||
curlv = cp.asnumpy(curl_v)
|
||||
curlw = cp.asnumpy(curl_w)
|
||||
|
||||
if verbose:
|
||||
print(' Done!', end ='')
|
||||
mtimer.end()
|
||||
|
||||
return (curlu, curlv, curlw)
|
||||
|
||||
def rotated_gradient(self, sfield, verbose=False):
|
||||
|
||||
if (sfield.shape != self.dims):
|
||||
#if (sfield.shape - self.dims).any():
|
||||
raise ValueError("Dimensions of scalar field should match that of the grid")
|
||||
|
||||
if self.dim != 2:
|
||||
raise ValueError("rotated_gradient works only for 2D")
|
||||
|
||||
if verbose:
|
||||
print(' Computing rotated gradient...', end ='')
|
||||
sys.stdout.flush()
|
||||
mtimer = Timer()
|
||||
|
||||
ddy, ddx = numpy.gradient(sfield, self.dx[0], self.dx[1])
|
||||
ddy *= -1.0
|
||||
|
||||
grad = numpy.stack((ddy, ddx), axis=-1)
|
||||
|
||||
if verbose:
|
||||
print(' Done!', end ='')
|
||||
mtimer.end()
|
||||
|
||||
return grad
|
||||
|
||||
def rotated_gradient_gpu(self, sfield, verbose=False):
|
||||
"""
|
||||
GPU-accelerated computation of rotated gradient for 2D scalar field.
|
||||
Uses CuPy for faster gradient computation on GPU.
|
||||
"""
|
||||
if not CUPY_AVAILABLE:
|
||||
if verbose:
|
||||
print('CuPy not available, falling back to CPU')
|
||||
return self.rotated_gradient(sfield, verbose)
|
||||
|
||||
if (sfield.shape != self.dims):
|
||||
raise ValueError("Dimensions of scalar field should match that of the grid")
|
||||
|
||||
if self.dim != 2:
|
||||
raise ValueError("rotated_gradient works only for 2D")
|
||||
|
||||
if verbose:
|
||||
print(' Computing rotated gradient on GPU...', end ='')
|
||||
sys.stdout.flush()
|
||||
mtimer = Timer()
|
||||
|
||||
# Transfer to GPU
|
||||
sfield_gpu = cp.asarray(sfield, dtype=cp.float32)
|
||||
|
||||
# Compute gradients on GPU
|
||||
ddy, ddx = cp.gradient(sfield_gpu, self.dx[0], self.dx[1])
|
||||
ddy *= -1.0
|
||||
|
||||
# Stack and transfer back to CPU
|
||||
grad_gpu = cp.stack((ddy, ddx), axis=-1)
|
||||
grad = cp.asnumpy(grad_gpu)
|
||||
|
||||
if verbose:
|
||||
print(' Done!', end ='')
|
||||
mtimer.end()
|
||||
|
||||
return grad
|
||||
|
||||
def gradient(self, sfield, verbose=False):
|
||||
|
||||
if (sfield.shape != self.dims):
|
||||
#if (sfield.shape - self.dims).any():
|
||||
raise ValueError("Dimensions of scalar field should match that of the grid")
|
||||
|
||||
if verbose:
|
||||
print(' Computing gradient...', end ='')
|
||||
sys.stdout.flush()
|
||||
mtimer = Timer()
|
||||
|
||||
if self.dim == 2:
|
||||
|
||||
# self.dx = (dy,dx)
|
||||
ddy, ddx = numpy.gradient(sfield, self.dx[0], self.dx[1])
|
||||
grad = numpy.stack((ddx, ddy), axis = -1)
|
||||
|
||||
elif self.dim == 3:
|
||||
|
||||
# self.dx = (dz,dy,dx)
|
||||
ddz, ddy, ddx = numpy.gradient(sfield, self.dx[0], self.dx[1], self.dx[2])
|
||||
grad = numpy.stack((ddx, ddy, ddz), axis = -1)
|
||||
|
||||
if verbose:
|
||||
print(' Done!', end ='')
|
||||
mtimer.end()
|
||||
|
||||
return grad
|
||||
|
||||
def gradient_gpu(self, sfield, verbose=False):
|
||||
"""
|
||||
GPU-accelerated computation of gradient for scalar field.
|
||||
Uses CuPy for faster gradient computation on GPU.
|
||||
"""
|
||||
if not CUPY_AVAILABLE:
|
||||
if verbose:
|
||||
print('CuPy not available, falling back to CPU')
|
||||
return self.gradient(sfield, verbose)
|
||||
|
||||
if (sfield.shape != self.dims):
|
||||
raise ValueError("Dimensions of scalar field should match that of the grid")
|
||||
|
||||
if verbose:
|
||||
print(' Computing gradient on GPU...', end ='')
|
||||
sys.stdout.flush()
|
||||
mtimer = Timer()
|
||||
|
||||
# Transfer to GPU
|
||||
sfield_gpu = cp.asarray(sfield, dtype=cp.float32)
|
||||
|
||||
if self.dim == 2:
|
||||
# Compute gradients on GPU
|
||||
ddy, ddx = cp.gradient(sfield_gpu, self.dx[0], self.dx[1])
|
||||
grad_gpu = cp.stack((ddx, ddy), axis=-1)
|
||||
|
||||
elif self.dim == 3:
|
||||
# Compute gradients on GPU
|
||||
ddz, ddy, ddx = cp.gradient(sfield_gpu, self.dx[0], self.dx[1], self.dx[2])
|
||||
grad_gpu = cp.stack((ddx, ddy, ddz), axis=-1)
|
||||
|
||||
# Transfer back to CPU
|
||||
grad = cp.asnumpy(grad_gpu)
|
||||
|
||||
if verbose:
|
||||
print(' Done!', end ='')
|
||||
mtimer.end()
|
||||
|
||||
return grad
|
||||
|
||||
def clear_gpu_cache(self):
|
||||
"""Clear GPU cache to free GPU memory"""
|
||||
if self._gpu_divcurl_cache is not None:
|
||||
self._gpu_divcurl_cache = None
|
||||
if CUPY_AVAILABLE:
|
||||
cp.get_default_memory_pool().free_all_blocks()
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# --------------------------------------------------------------------------
|
||||
@@ -0,0 +1,16 @@
|
||||
"""
|
||||
Low-level kernels and shared numeric helpers for HHD.
|
||||
|
||||
Use this file for pieces that should not live inside the public `nHHD` API or the
|
||||
`PoissonSolver` class, for example:
|
||||
- finite-difference helper functions
|
||||
- small backend-dispatch helpers
|
||||
- reusable CuPy/OpenCL kernel preparation code
|
||||
- shared CPU math that does not depend on object state
|
||||
"""
|
||||
|
||||
# No standalone helpers from the old monolithic `HHD.py` are required yet.
|
||||
# Keep this module in place so future extracted CUDA/OpenCL/CPU helpers have a
|
||||
# stable home without changing the public package layout again.
|
||||
|
||||
__all__ = []
|
||||
@@ -0,0 +1,954 @@
|
||||
import numpy
|
||||
import timeit
|
||||
import sys
|
||||
# import cv2
|
||||
from scipy import signal, spatial
|
||||
|
||||
from ..backend_availability import (
|
||||
CUPY_AVAILABLE, OPENCL_AVAILABLE, KERNEL_AVAILABLE,
|
||||
cp, cp_signal, _torque_kernel,
|
||||
cl, cla, cl_rfftn, cl_irfftn
|
||||
)
|
||||
from .timing import Timer
|
||||
|
||||
'''
|
||||
# ------------------------------------------------------------------------------
|
||||
Poisson Solver for 1D, 2D, and 3D scalar fields.
|
||||
Data may be defined on a rectlinear grid or unstructured points.
|
||||
|
||||
1) create an object
|
||||
PoissonSolver(args)
|
||||
'solver' : type of solver
|
||||
'S' for spatial convolution
|
||||
only option for unstructured points
|
||||
'F' for convolution in frequency space
|
||||
default option for regular grids
|
||||
|
||||
--> for regular grid
|
||||
'grid' : shape of the rectlinear grid
|
||||
tuple of dimensions (Z,Y,X)
|
||||
'spacings' : spacings of the rectlinear grid
|
||||
tuple of spacings (dz,dy,dx)
|
||||
|
||||
--> for unstrucutred points
|
||||
'points' : unstructured points
|
||||
numpy array of shape (N, dim)
|
||||
'pvolumes' : point volumes for each point
|
||||
numpy array of shape (N, 1)
|
||||
|
||||
2) call prepare()
|
||||
this initializes internal datastructures depending upon the mesh.
|
||||
|
||||
Once initialized, solve() can be called any number of times to
|
||||
solve the Poisson equation on the same domain. If the domain
|
||||
changes (e.g., a Lagrangian mesh), create a new object.
|
||||
|
||||
3) call solve(f), where, f is a numpy array defining a scalar field
|
||||
|
||||
for rectlinear grid, f.shape = grid.shape
|
||||
for unstructured points, f.shape = (points.shape[0], 1)
|
||||
|
||||
# ------------------------------------------------------------------------------
|
||||
'''
|
||||
|
||||
class PoissonSolver(object):
|
||||
|
||||
# constructor
|
||||
def __init__(self, **kwargs):
|
||||
|
||||
args = kwargs.keys()
|
||||
|
||||
if ('grid' in args) == ('points' in args):
|
||||
raise SyntaxError("Poisson Solver needs either shape of a regular grid or the points in an rectilinear/unstructured grid")
|
||||
|
||||
if ('grid' in args) != ('spacings' in args):
|
||||
raise SyntaxError("Poisson Solver needs spacings for regular grid")
|
||||
|
||||
if ('points' in args) != ('pvolumes' in args):
|
||||
raise SyntaxError("Poisson Solver needs point volumes for unstructured points")
|
||||
|
||||
# (default) data type
|
||||
self.dtype = numpy.float32
|
||||
if 'dtype' in args:
|
||||
self.dtype = kwargs['dtype']
|
||||
|
||||
if 'grid' in args:
|
||||
self.ptype = 'G'
|
||||
self.stype = 'F' # default
|
||||
self.gdims = kwargs['grid']
|
||||
self.gdx = kwargs['spacings']
|
||||
self.dim = len(self.gdims)
|
||||
if self.dim != len(self.gdx):
|
||||
raise ValueError("Dimensions of spacings should match that of the grid")
|
||||
|
||||
elif 'points' in args:
|
||||
self.ptype = 'P'
|
||||
self.stype = 'S' # only option
|
||||
self.points = kwargs['points']
|
||||
self.pvolumes = kwargs['pvolumes']
|
||||
self.dim = len(self.points.shape)
|
||||
if self.pvolumes.shape[0] != self.points.shape[0]:
|
||||
raise ValueError("Number of pvolumes should match that of the points")
|
||||
|
||||
if 'solver' in args:
|
||||
if (kwargs['solver'] != 'S' and kwargs['solver'] != 'F'):
|
||||
raise SyntaxError('Solver can be only of type S(patial) or F(requency)')
|
||||
|
||||
if (kwargs['solver'] == 'F') and ('points' in args):
|
||||
raise SyntaxError('Only Spatial solver supported for unstructured points')
|
||||
|
||||
if (kwargs['solver'] == 'S') and ('grid' in args):
|
||||
raise SyntaxError('Frequncy solver strongly suggested for structured grids')
|
||||
|
||||
self.stype = kwargs['solver']
|
||||
if 'poisson_backend' in args:
|
||||
self.poisson_backend = kwargs['poisson_backend']
|
||||
if self.poisson_backend not in ('cpu', 'opencl', 'cupy'):
|
||||
raise ValueError("poisson_backend must be one of: 'cpu', 'opencl', 'cupy'")
|
||||
|
||||
if (self.dim > 3):
|
||||
raise ValueError("Poisson Solver works only for 1, 2, or 3 dimensions")
|
||||
|
||||
self.ready = False
|
||||
if self.poisson_backend == 'cupy':
|
||||
self._gpu_cache = None # Cache for GPU arrays
|
||||
if self.poisson_backend == 'opencl':
|
||||
self._opencl_cache = None # Cache for OpenCL FFT buffers
|
||||
self._opencl_queue = None # OpenCL command queue
|
||||
self._opencl_context = None # OpenCL context (set when queue is created)
|
||||
self._opencl_crop_kernel = None # Cached OpenCL crop kernel
|
||||
self._opencl_scatter_kernel = None # Cached scatter (reverse-crop) kernel
|
||||
self._opencl_fused_mul_kernel = None # Cached fused complex multiply kernel
|
||||
self._opencl_scatter_batch_kernel = None # Cached batch scatter kernel
|
||||
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'):
|
||||
# print( 'PoissonSolver', self.dim, 'D points =', self.points.shape)
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# create the Green's function in appropriate representation
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
# generate a grid -(N-1) to (N-1) to define distance function
|
||||
# where N is the size of the original grid
|
||||
def _create_radialGrid(self, verbose):
|
||||
|
||||
mtimer = Timer()
|
||||
if verbose:
|
||||
print( ' - creating distance kernel:', end ='')
|
||||
|
||||
rdims = tuple(2*d - 1 for d in self.gdims)
|
||||
|
||||
R = numpy.indices(rdims, dtype=self.dtype)
|
||||
for d in range(self.dim):
|
||||
R[d] += 1-self.gdims[d]
|
||||
R[d] *= self.gdx[d]
|
||||
|
||||
if self.dim == 1:
|
||||
numpy.absolute(R[0], R[0])
|
||||
R = R[0]
|
||||
|
||||
elif self.dim == 2:
|
||||
numpy.hypot(R[1], R[0], R[0])
|
||||
R = R[0]
|
||||
|
||||
# half of the average grid spacing
|
||||
zval = 0.5 * (self.gdx[0]+self.gdx[1]) / 2.0
|
||||
R[self.gdims[0]-1, self.gdims[1]-1] = zval
|
||||
|
||||
elif self.dim == 3:
|
||||
numpy.hypot(R[2], R[1], R[1])
|
||||
numpy.hypot(R[1], R[0], R[0])
|
||||
R = R[0]
|
||||
|
||||
# half of the average grid spacing
|
||||
zval = 0.5 * (self.gdx[0]+self.gdx[1]+self.gdx[2]) / 3.0
|
||||
R[self.gdims[0]-1, self.gdims[1]-1, self.gdims[2]-1] = zval
|
||||
|
||||
if verbose:
|
||||
print( R.shape, R.min(), R.max())
|
||||
mtimer.end()
|
||||
|
||||
return R
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# scale a function by point volumes for correct integration
|
||||
def _scale(self, func):
|
||||
|
||||
if self.ptype == 'P':
|
||||
return numpy.multiply(func, self.pvolumes)
|
||||
|
||||
elif self.ptype == 'G':
|
||||
sfunc = func * numpy.prod(self.gdx)
|
||||
|
||||
'''
|
||||
# volume of the boundary elements is only half the voxel
|
||||
if self.dim == 1:
|
||||
sfunc[ 0] *= 0.5
|
||||
sfunc[-1] *= 0.5
|
||||
elif self.dim == 2:
|
||||
sfunc[ 0, :] *= 0.5
|
||||
sfunc[-1, :] *= 0.5
|
||||
sfunc[ :, 0] *= 0.5
|
||||
sfunc[ :,-1] *= 0.5
|
||||
elif self.dim == 3:
|
||||
sfunc[ 0, :, :] *= 0.5
|
||||
sfunc[-1, :, :] *= 0.5
|
||||
sfunc[ :, 0, :] *= 0.5
|
||||
sfunc[ :,-1, :] *= 0.5
|
||||
sfunc[ :, :, 0] *= 0.5
|
||||
sfunc[ :, :,-1] *= 0.5
|
||||
'''
|
||||
return sfunc
|
||||
|
||||
def prepare(self, verbose=False):
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
gtimer = Timer()
|
||||
if verbose:
|
||||
print( '\nInitializing Poisson solver, type =', self.stype)
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# generate a grid to store pairwise distances to store Green's function
|
||||
self.G = None # definition depends upon solver type
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# for Fourier solution,
|
||||
if self.stype == 'F':
|
||||
|
||||
if self.ptype == 'G':
|
||||
self.G = self._create_radialGrid(verbose)
|
||||
|
||||
elif self.ptype == 'P':
|
||||
raise SyntaxError('Frequncy solver not supported for unstructured points')
|
||||
#self.G = create_distArray()
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# for spatial solution,
|
||||
# use the pairwise distance matrix, N x N
|
||||
# where N is the number of points
|
||||
elif self.stype == 'S':
|
||||
|
||||
if self.ptype == 'G':
|
||||
raise SyntaxError('Spatial solver not suggested for structured grids')
|
||||
#self.G = self.create_points(verbose)
|
||||
|
||||
elif self.ptype == 'P':
|
||||
# for 1D points, add an extra axis
|
||||
if (self.dim == 1) and (len(self.points.shape) == 1):
|
||||
self.points = self.points[:,numpy.newaxis]
|
||||
|
||||
# create pairwise distance matrix now
|
||||
ltimer = Timer()
|
||||
self.G = spatial.distance.cdist(self.points, self.points)
|
||||
|
||||
# TODO: use a good approximation here
|
||||
zval = numpy.power(10.0, -10.0)
|
||||
self.G += zval * numpy.identity(self.G.shape[0])
|
||||
|
||||
if verbose:
|
||||
print(' - created pairwise distance matrix:', self.G.shape, self.G.min(), self.G.max(), end = "")
|
||||
ltimer.end()
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# compute the Green's function
|
||||
ltimer = Timer()
|
||||
if verbose:
|
||||
print( ' - computing the Green\'s function:')
|
||||
|
||||
if self.dim == 1:
|
||||
numpy.multiply(self.G, 0.5, self.G)
|
||||
|
||||
elif self.dim == 2:
|
||||
numpy.log(self.G, self.G)
|
||||
numpy.multiply(self.G, (0.5 / numpy.pi), self.G)
|
||||
|
||||
elif self.dim == 3:
|
||||
numpy.reciprocal(self.G, self.G)
|
||||
numpy.multiply(self.G, (-0.25 / numpy.pi), self.G)
|
||||
|
||||
if verbose:
|
||||
print( self.G.shape, self.G.min(), self.G.max(), end = "")
|
||||
ltimer.end()
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
self.ready = True
|
||||
|
||||
if verbose:
|
||||
print( 'Poisson solver initialized', end ='')
|
||||
gtimer.end()
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# compute the integral solution
|
||||
# --------------------------------------------------------------------------
|
||||
def solve(self, f, verbose=False):
|
||||
|
||||
if not self.ready:
|
||||
self.prepare(verbose)
|
||||
|
||||
fshape = f.shape
|
||||
|
||||
gtimer = Timer()
|
||||
if verbose:
|
||||
print( '\nSolving Poisson Eq.')
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# regular grid
|
||||
if self.ptype == 'G':
|
||||
|
||||
#if (self.gdims - fshape).any():
|
||||
if self.gdims != fshape:
|
||||
print( self.gdims, fshape)
|
||||
raise ValueError("Shape of function should match shape of grid")
|
||||
|
||||
# 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)
|
||||
|
||||
# # 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)
|
||||
# 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
|
||||
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':
|
||||
raise SyntaxError('Spatial solver not suggested for structured grids')
|
||||
#p = (self.G).dot(f.reshape(-1)).reshape(fshape)
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# unstructured points
|
||||
elif self.ptype == 'P':
|
||||
|
||||
if (self.points.shape[0] != fshape[0]) or (1 != len(fshape)):
|
||||
raise ValueError("Shape of function should be (N,1) with N as the number of points")
|
||||
|
||||
# convolution in spatial domain
|
||||
# Dec 13, 2017 -- pre scaling works for unstructured grids!
|
||||
if self.stype == 'S':
|
||||
p = numpy.multiply(f, self.pvolumes)
|
||||
numpy.dot(self.G, p, p)
|
||||
#numpy.multiply(p, self.pvolumes, p)
|
||||
|
||||
elif self.stype == 'F':
|
||||
raise SyntaxError('Frequncy solver not supported for unstructured points')
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
if verbose:
|
||||
print( 'Poisson solver finished')
|
||||
gtimer.end()
|
||||
|
||||
return p
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# GPU-accelerated solve using CuPy
|
||||
# --------------------------------------------------------------------------
|
||||
def solve_gpu(self, f, verbose=False, return_gpu=False):
|
||||
"""
|
||||
GPU-accelerated Poisson solver using CuPy.
|
||||
Falls back to CPU if CuPy is not available.
|
||||
|
||||
Args:
|
||||
f: Input field (numpy array)
|
||||
verbose: Print timing information
|
||||
|
||||
Returns:
|
||||
p: Solution (numpy array on CPU)
|
||||
"""
|
||||
if not CUPY_AVAILABLE:
|
||||
if verbose:
|
||||
print("CuPy not available, falling back to CPU solver")
|
||||
return self.solve(f, verbose)
|
||||
|
||||
if not self.ready:
|
||||
self.prepare(verbose)
|
||||
|
||||
fshape = f.shape
|
||||
|
||||
gtimer = Timer()
|
||||
if verbose:
|
||||
print('\nSolving Poisson Eq. on GPU')
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# regular grid
|
||||
if self.ptype == 'G':
|
||||
if self.gdims != fshape:
|
||||
print(self.gdims, fshape)
|
||||
raise ValueError("Shape of function should match shape of grid")
|
||||
|
||||
# convolution in frequency domain using GPU
|
||||
if self.stype == 'F':
|
||||
# Initialize GPU cache if needed
|
||||
if self._gpu_cache is None:
|
||||
if verbose:
|
||||
print(' - Transferring Green\'s function to GPU...')
|
||||
self._gpu_cache = {
|
||||
'G_gpu': cp.asarray(self.G, dtype=cp.float32),
|
||||
'gdx_prod': numpy.prod(self.gdx)
|
||||
}
|
||||
|
||||
# Transfer input to GPU
|
||||
f_gpu = cp.asarray(f, dtype=cp.float32)
|
||||
|
||||
# Perform FFT convolution on GPU
|
||||
# CuPy's fftconvolve is GPU-accelerated
|
||||
p_gpu = cp_signal.fftconvolve(f_gpu, self._gpu_cache['G_gpu'], mode='same')
|
||||
|
||||
# Scale by grid spacing (on GPU)
|
||||
p_gpu *= self._gpu_cache['gdx_prod']
|
||||
|
||||
if verbose:
|
||||
print('Poisson solver (GPU) finished')
|
||||
gtimer.end()
|
||||
|
||||
# Transfer result back to CPU
|
||||
if return_gpu:
|
||||
return p_gpu # Return CuPy array
|
||||
return cp.asnumpy(p_gpu)
|
||||
# p = cp.asnumpy(p_gpu)
|
||||
|
||||
elif self.stype == 'S':
|
||||
raise SyntaxError('Spatial solver not suggested for structured grids')
|
||||
|
||||
def _solve_gpu_internal(self, f_gpu):
|
||||
"""Internal GPU solve that accepts and returns GPU arrays."""
|
||||
"""
|
||||
GPU-accelerated Poisson solver using CuPy.
|
||||
Falls back to CPU if CuPy is not available.
|
||||
|
||||
Args:
|
||||
f_gpu: Input field (numpy array)
|
||||
|
||||
Returns:
|
||||
p_gpu: Solution (CuPy array)
|
||||
"""
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# regular grid
|
||||
if self.ptype == 'G':
|
||||
|
||||
# convolution in frequency domain using GPU
|
||||
if self.stype == 'F':
|
||||
# Initialize GPU cache if needed
|
||||
if self._gpu_cache is None:
|
||||
self._gpu_cache = {
|
||||
'G_gpu': cp.asarray(self.G, dtype=cp.float32),
|
||||
'gdx_prod': numpy.prod(self.gdx)
|
||||
}
|
||||
|
||||
# Perform FFT convolution on GPU
|
||||
# CuPy's fftconvolve is GPU-accelerated
|
||||
p_gpu = cp_signal.fftconvolve(f_gpu, self._gpu_cache['G_gpu'], mode='same')
|
||||
|
||||
# Scale by grid spacing (on GPU)
|
||||
p_gpu *= self._gpu_cache['gdx_prod']
|
||||
|
||||
return p_gpu # Return CuPy array
|
||||
|
||||
def clear_gpu_cache(self):
|
||||
"""Clear GPU cache to free GPU memory"""
|
||||
if self._gpu_cache is not None:
|
||||
self._gpu_cache = None
|
||||
if CUPY_AVAILABLE:
|
||||
cp.get_default_memory_pool().free_all_blocks()
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# unstructured points - not optimized for GPU
|
||||
# elif self.ptype == 'P':
|
||||
# if verbose:
|
||||
# print(' - GPU solver not optimized for unstructured points, using CPU')
|
||||
# return self.solve(f, verbose)
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
|
||||
# return p
|
||||
|
||||
def _get_opencl_queue(self):
|
||||
"""Create and cache an OpenCL queue on first available GPU."""
|
||||
if self._opencl_queue is not None:
|
||||
return self._opencl_queue
|
||||
|
||||
if not OPENCL_AVAILABLE:
|
||||
raise RuntimeError("OpenCL/pyvkfft is not available")
|
||||
|
||||
ctx = None
|
||||
for platform in cl.get_platforms():
|
||||
for device in platform.get_devices():
|
||||
if device.type & cl.device_type.GPU == 0:
|
||||
continue
|
||||
ctx = cl.Context(devices=(device,))
|
||||
break
|
||||
if ctx is not None:
|
||||
break
|
||||
|
||||
if ctx is None:
|
||||
raise RuntimeError("No OpenCL GPU device found")
|
||||
|
||||
self._opencl_queue = cl.CommandQueue(ctx)
|
||||
self._opencl_context = ctx
|
||||
return self._opencl_queue
|
||||
|
||||
def _get_opencl_crop_kernel(self, queue):
|
||||
"""Create and cache OpenCL kernel for cropping convolution result."""
|
||||
if hasattr(self, '_opencl_crop_kernel') and self._opencl_crop_kernel is not None:
|
||||
return self._opencl_crop_kernel
|
||||
|
||||
# Simple 2D copy kernel
|
||||
kernel_code = """
|
||||
__kernel void crop_and_scale(
|
||||
__global const float* src,
|
||||
__global float* dst,
|
||||
const int src_width,
|
||||
const int dst_width,
|
||||
const int dst_height,
|
||||
const int start_x,
|
||||
const int start_y,
|
||||
const float scale
|
||||
) {
|
||||
int gid_x = get_global_id(0);
|
||||
int gid_y = get_global_id(1);
|
||||
|
||||
if (gid_x < dst_width && gid_y < dst_height) {
|
||||
int src_idx = (start_y + gid_y) * src_width + (start_x + gid_x);
|
||||
int dst_idx = gid_y * dst_width + gid_x;
|
||||
dst[dst_idx] = src[src_idx] * scale;
|
||||
}
|
||||
}
|
||||
"""
|
||||
|
||||
program = cl.Program(self._opencl_context, kernel_code).build()
|
||||
self._opencl_crop_kernel = program.crop_and_scale
|
||||
return self._opencl_crop_kernel
|
||||
|
||||
def _get_opencl_scatter_kernel(self, queue):
|
||||
"""Kernel that copies unpadded input into a zero-padded device buffer (reverse of crop)."""
|
||||
if self._opencl_scatter_kernel is not None:
|
||||
return self._opencl_scatter_kernel
|
||||
|
||||
kernel_code = """
|
||||
__kernel void scatter_to_padded(
|
||||
__global const float* src,
|
||||
__global float* dst,
|
||||
const int src_width,
|
||||
const int src_height,
|
||||
const int dst_width
|
||||
) {
|
||||
int gid_x = get_global_id(0);
|
||||
int gid_y = get_global_id(1);
|
||||
if (gid_x < src_width && gid_y < src_height) {
|
||||
dst[gid_y * dst_width + gid_x] = src[gid_y * src_width + gid_x];
|
||||
}
|
||||
}
|
||||
"""
|
||||
program = cl.Program(self._opencl_context, kernel_code).build()
|
||||
self._opencl_scatter_kernel = program.scatter_to_padded
|
||||
return self._opencl_scatter_kernel
|
||||
|
||||
def _get_opencl_fused_mul_kernel(self, queue):
|
||||
"""Kernel that multiplies two complex spectra by one Green's spectrum in a single pass."""
|
||||
if self._opencl_fused_mul_kernel is not None:
|
||||
return self._opencl_fused_mul_kernel
|
||||
|
||||
kernel_code = """
|
||||
__kernel void fused_complex_mul(
|
||||
__global float2* fft1,
|
||||
__global float2* fft2,
|
||||
__global const float2* g,
|
||||
const int n
|
||||
) {
|
||||
int i = get_global_id(0);
|
||||
if (i < n) {
|
||||
float2 gv = g[i];
|
||||
float2 a = fft1[i];
|
||||
float2 b = fft2[i];
|
||||
fft1[i] = (float2)(a.x * gv.x - a.y * gv.y,
|
||||
a.x * gv.y + a.y * gv.x);
|
||||
fft2[i] = (float2)(b.x * gv.x - b.y * gv.y,
|
||||
b.x * gv.y + b.y * gv.x);
|
||||
}
|
||||
}
|
||||
"""
|
||||
program = cl.Program(self._opencl_context, kernel_code).build()
|
||||
self._opencl_fused_mul_kernel = program.fused_complex_mul
|
||||
return self._opencl_fused_mul_kernel
|
||||
|
||||
def _get_opencl_scatter_batch_kernel(self, queue):
|
||||
"""Batched scatter: copies two unpadded inputs into a stacked zero-padded device buffer."""
|
||||
if self._opencl_scatter_batch_kernel is not None:
|
||||
return self._opencl_scatter_batch_kernel
|
||||
|
||||
kernel_code = """
|
||||
__kernel void scatter_batch(
|
||||
__global const float* src,
|
||||
__global float* dst,
|
||||
const int src_width,
|
||||
const int src_height,
|
||||
const int dst_width,
|
||||
const int src_plane,
|
||||
const int dst_plane
|
||||
) {
|
||||
int gid_x = get_global_id(0);
|
||||
int gid_y = get_global_id(1);
|
||||
int gid_z = get_global_id(2);
|
||||
if (gid_x < src_width && gid_y < src_height) {
|
||||
int src_idx = gid_z * src_plane + gid_y * src_width + gid_x;
|
||||
int dst_idx = gid_z * dst_plane + gid_y * dst_width + gid_x;
|
||||
dst[dst_idx] = src[src_idx];
|
||||
}
|
||||
}
|
||||
"""
|
||||
program = cl.Program(self._opencl_context, kernel_code).build()
|
||||
self._opencl_scatter_batch_kernel = program.scatter_batch
|
||||
return self._opencl_scatter_batch_kernel
|
||||
|
||||
def _get_opencl_crop_batch_kernel(self, queue):
|
||||
"""Batched crop+scale: extracts 'same' region from two stacked full-conv results."""
|
||||
if self._opencl_crop_batch_kernel is not None:
|
||||
return self._opencl_crop_batch_kernel
|
||||
|
||||
kernel_code = """
|
||||
__kernel void crop_batch(
|
||||
__global const float* src,
|
||||
__global float* dst,
|
||||
const int src_width,
|
||||
const int dst_width,
|
||||
const int dst_height,
|
||||
const int start_x,
|
||||
const int start_y,
|
||||
const float scale,
|
||||
const int src_plane,
|
||||
const int dst_plane
|
||||
) {
|
||||
int gid_x = get_global_id(0);
|
||||
int gid_y = get_global_id(1);
|
||||
int gid_z = get_global_id(2);
|
||||
if (gid_x < dst_width && gid_y < dst_height) {
|
||||
int src_idx = gid_z * src_plane + (start_y + gid_y) * src_width + (start_x + gid_x);
|
||||
int dst_idx = gid_z * dst_plane + gid_y * dst_width + gid_x;
|
||||
dst[dst_idx] = src[src_idx] * scale;
|
||||
}
|
||||
}
|
||||
"""
|
||||
program = cl.Program(self._opencl_context, kernel_code).build()
|
||||
self._opencl_crop_batch_kernel = program.crop_batch
|
||||
return self._opencl_crop_batch_kernel
|
||||
|
||||
def _get_opencl_batch_mul_kernel(self, queue):
|
||||
"""Batched complex multiply: multiplies each batch plane by the same Green's spectrum.
|
||||
Uses multiple elements per work item, __constant for broadcast-read Green's spectrum,
|
||||
and fma for better throughput/precision."""
|
||||
if self._opencl_batch_mul_kernel is not None:
|
||||
return self._opencl_batch_mul_kernel
|
||||
|
||||
kernel_code = """
|
||||
__kernel void batch_complex_mul(
|
||||
__global float2* fft_batch,
|
||||
__global const float2* g,
|
||||
const int n
|
||||
) {
|
||||
const int b = get_global_id(0);
|
||||
const int base = b * n;
|
||||
const int i0 = get_global_id(1) * 4;
|
||||
const int i_end = min(i0 + 4, n);
|
||||
for (int i = i0; i < i_end; i++) {
|
||||
float2 gv = g[i];
|
||||
float2 a = fft_batch[base + i];
|
||||
fft_batch[base + i] = (float2)(fma(a.x, gv.x, -(a.y * gv.y)),
|
||||
fma(a.x, gv.y, a.y * gv.x));
|
||||
}
|
||||
}
|
||||
"""
|
||||
program = cl.Program(self._opencl_context, kernel_code).build()
|
||||
self._opencl_batch_mul_kernel = program.batch_complex_mul
|
||||
return self._opencl_batch_mul_kernel
|
||||
|
||||
def _get_opencl_alloc_host_ptr(self):
|
||||
"""Return an allocator that creates buffers with ALLOC_HOST_PTR for pinned host access (reduces H2D/D2H transfer cost)."""
|
||||
ctx = self._opencl_context
|
||||
mf = cl.mem_flags
|
||||
|
||||
def allocator(nbytes):
|
||||
return cl.Buffer(ctx, mf.READ_WRITE | mf.ALLOC_HOST_PTR, size=nbytes)
|
||||
|
||||
return allocator
|
||||
|
||||
@staticmethod
|
||||
def _center_crop_same(full_conv, target_shape):
|
||||
"""Crop full convolution result to scipy 'same' output shape."""
|
||||
starts = [(full_conv.shape[i] - target_shape[i]) // 2 for i in range(len(target_shape))]
|
||||
slices = tuple(slice(starts[i], starts[i] + target_shape[i]) for i in range(len(target_shape)))
|
||||
return full_conv[slices]
|
||||
|
||||
|
||||
def solve_opencl_batch2(self, f1, f2, verbose=False, profile=False):
|
||||
"""
|
||||
Batched OpenCL Poisson solver for two 2D fields.
|
||||
Optimized: scatter H2D (~4x smaller transfer), fused multiply
|
||||
(reads g_fft once), async D2H (one sync instead of two).
|
||||
When profile=True, records and optionally prints per-stage timings.
|
||||
"""
|
||||
if not OPENCL_AVAILABLE:
|
||||
if verbose:
|
||||
print("OpenCL/pyvkfft not available, falling back to CPU solver")
|
||||
return self.solve(f1, verbose), self.solve(f2, verbose)
|
||||
|
||||
# if self.dim != 2:
|
||||
# return self.solve_opencl(f1, verbose, profile), self.solve_opencl(f2, verbose, profile)
|
||||
|
||||
if not self.ready:
|
||||
self.prepare(verbose)
|
||||
|
||||
fshape = f1.shape
|
||||
if fshape != f2.shape:
|
||||
raise ValueError("solve_opencl_batch2 requires inputs with identical shapes")
|
||||
|
||||
if self.ptype != 'G' or self.stype != 'F':
|
||||
return self.solve_opencl(f1, verbose, profile), self.solve_opencl(f2, verbose, profile)
|
||||
|
||||
if self.gdims != fshape:
|
||||
print(self.gdims, fshape)
|
||||
raise ValueError("Shape of function should match shape of grid")
|
||||
|
||||
timing = {'cache_prepare_ms': 0.0}
|
||||
t_total = timeit.default_timer()
|
||||
|
||||
t0 = timeit.default_timer()
|
||||
queue = self._get_opencl_queue()
|
||||
timing['queue_ms'] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
t0 = timeit.default_timer()
|
||||
fshape = f1.shape
|
||||
full_shape = tuple(2 * s - 1 for s in fshape)
|
||||
timing['host_prepare_ms'] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
cache_valid = (
|
||||
self._opencl_cache is not None and
|
||||
self._opencl_cache['fshape'] == fshape and
|
||||
self._opencl_cache['full_shape'] == full_shape and
|
||||
'f_batch_dev' in self._opencl_cache
|
||||
)
|
||||
|
||||
if not cache_valid:
|
||||
t0 = timeit.default_timer()
|
||||
g_host = numpy.asarray(self.G, dtype=numpy.float32)
|
||||
g_pad = numpy.zeros(full_shape, dtype=numpy.float32)
|
||||
g_slices = tuple(slice(0, s) for s in g_host.shape)
|
||||
g_pad[g_slices] = g_host
|
||||
g_dev = cla.to_device(queue, g_pad)
|
||||
g_fft_dev = cl_rfftn(g_dev)
|
||||
if profile:
|
||||
queue.finish()
|
||||
|
||||
# For circular conv at size (2s-1) with G of shape (2s-1) and f of shape s,
|
||||
# alias-free positions are [s-1, 2s-2], matching scipy 'same' extraction.
|
||||
same_starts = [fshape[i] - 1 for i in range(self.dim)]
|
||||
fft_shape = tuple(list(full_shape[:-1]) + [full_shape[-1] // 2 + 1])
|
||||
fft_n_elements = 1
|
||||
for d in fft_shape:
|
||||
fft_n_elements *= d
|
||||
src_plane = fshape[0] * fshape[1]
|
||||
dst_plane = full_shape[0] * full_shape[1]
|
||||
|
||||
# Pinned allocator for staging (H2D) and output (D2H) buffers
|
||||
alloc_pinned = self._get_opencl_alloc_host_ptr()
|
||||
|
||||
self._opencl_cache = {
|
||||
'fshape': fshape,
|
||||
'full_shape': full_shape,
|
||||
'fft_shape': fft_shape,
|
||||
'fft_n_elements': fft_n_elements,
|
||||
'g_fft_dev': g_fft_dev,
|
||||
'same_starts': same_starts,
|
||||
'gdx_prod': numpy.prod(self.gdx),
|
||||
'src_plane': src_plane,
|
||||
'dst_plane': dst_plane,
|
||||
'f_batch_dev': cla.zeros(queue, (2,) + full_shape, dtype=numpy.float32),
|
||||
'host_stage': numpy.empty((2,) + fshape, dtype=numpy.float32),
|
||||
'f_stage_dev': cla.empty(queue, (2,) + fshape, dtype=numpy.float32),
|
||||
'f_fft_batch_dev': cla.empty(queue, (2,) + fft_shape, dtype=numpy.complex64),
|
||||
'conv_batch_dev': cla.empty(queue, (2,) + full_shape, dtype=numpy.float32),
|
||||
'p_batch_dev': cla.empty(queue, (2,) + fshape, dtype=numpy.float32, allocator=alloc_pinned),
|
||||
'p1_host': numpy.empty(fshape, dtype=numpy.float32),
|
||||
'p2_host': numpy.empty(fshape, dtype=numpy.float32),
|
||||
}
|
||||
timing['cache_prepare_ms'] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
c = self._opencl_cache
|
||||
fshape_c = c['fshape']
|
||||
full_shape_c = c['full_shape']
|
||||
|
||||
# H2D: compact transfer of unpadded data, then scatter to padded buffer on device
|
||||
t0 = timeit.default_timer()
|
||||
H, W = fshape_c
|
||||
f1_f32 = f1 if (f1.flags['C_CONTIGUOUS'] and f1.dtype == numpy.float32) else numpy.ascontiguousarray(f1, dtype=numpy.float32)
|
||||
f2_f32 = f2 if (f2.flags['C_CONTIGUOUS'] and f2.dtype == numpy.float32) else numpy.ascontiguousarray(f2, dtype=numpy.float32)
|
||||
host_stage = c['host_stage']
|
||||
numpy.copyto(host_stage[0], f1_f32)
|
||||
numpy.copyto(host_stage[1], f2_f32)
|
||||
f_stage_dev = c['f_stage_dev']
|
||||
cl.enqueue_copy(queue, f_stage_dev.data, host_stage, is_blocking=False)
|
||||
c['f_batch_dev'].fill(0, queue=queue)
|
||||
scatter_batch = self._get_opencl_scatter_batch_kernel(queue)
|
||||
scatter_batch(
|
||||
queue, (W, H, 2), None,
|
||||
f_stage_dev.data, c['f_batch_dev'].data,
|
||||
numpy.int32(W), numpy.int32(H),
|
||||
numpy.int32(full_shape_c[1]),
|
||||
numpy.int32(H * W),
|
||||
numpy.int32(c['dst_plane']),
|
||||
)
|
||||
if profile:
|
||||
queue.finish()
|
||||
timing['h2d_set_ms'] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
# Single batched 2D FFT for both fields (ndim=2 treats dim-0 as batch)
|
||||
t0 = timeit.default_timer()
|
||||
fft_batch = cl_rfftn(c['f_batch_dev'], c['f_fft_batch_dev'], ndim=2)
|
||||
if profile:
|
||||
queue.finish()
|
||||
timing['fft_f_ms'] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
# Batched multiply: each batch plane * same Green's spectrum (4 elements/work item, explicit local size)
|
||||
t0 = timeit.default_timer()
|
||||
n_fft = c['fft_n_elements']
|
||||
global_n = max(1, (n_fft + 3) // 4)
|
||||
local_n = min(256, global_n)
|
||||
global_n = ((global_n + local_n - 1) // local_n) * local_n
|
||||
batch_mul = self._get_opencl_batch_mul_kernel(queue)
|
||||
batch_mul(
|
||||
queue, (2, global_n), (1, local_n),
|
||||
fft_batch.data,
|
||||
c['g_fft_dev'].data,
|
||||
numpy.int32(n_fft)
|
||||
)
|
||||
if profile:
|
||||
queue.finish()
|
||||
timing['mul_ms'] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
# Single batched 2D iFFT for both fields
|
||||
t0 = timeit.default_timer()
|
||||
conv_batch = cl_irfftn(fft_batch, c['conv_batch_dev'], ndim=2)
|
||||
if profile:
|
||||
queue.finish()
|
||||
timing['ifft_ms'] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
# Batch crop + scale on device
|
||||
t0 = timeit.default_timer()
|
||||
crop_batch = self._get_opencl_crop_batch_kernel(queue)
|
||||
same_starts = c['same_starts']
|
||||
crop_batch(
|
||||
queue, (fshape_c[1], fshape_c[0], 2), None,
|
||||
conv_batch.data, c['p_batch_dev'].data,
|
||||
numpy.int32(full_shape_c[1]),
|
||||
numpy.int32(fshape_c[1]),
|
||||
numpy.int32(fshape_c[0]),
|
||||
numpy.int32(same_starts[1]),
|
||||
numpy.int32(same_starts[0]),
|
||||
numpy.float32(c['gdx_prod']),
|
||||
numpy.int32(c['dst_plane']),
|
||||
numpy.int32(c['src_plane'])
|
||||
)
|
||||
if profile:
|
||||
queue.finish()
|
||||
timing['crop_scale_ms'] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
# D2H: map pinned buffer, copy into pre-allocated host arrays (zero allocation)
|
||||
t0 = timeit.default_timer()
|
||||
p1 = c['p1_host']
|
||||
p2 = c['p2_host']
|
||||
mapped_p = c['p_batch_dev'].map_to_host(queue=queue, flags=cl.map_flags.READ, is_blocking=True)
|
||||
numpy.copyto(p1, mapped_p[0])
|
||||
numpy.copyto(p2, mapped_p[1])
|
||||
mapped_p.base.release(queue=queue)
|
||||
timing['d2h_get_ms'] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
timing['total_ms'] = (timeit.default_timer() - t_total) * 1000.0
|
||||
self._opencl_last_timing = timing
|
||||
|
||||
if verbose:
|
||||
print('Poisson solver (OpenCL FFT batch2) finished')
|
||||
if profile:
|
||||
print(
|
||||
"OpenCL batch2 timing [ms]: "
|
||||
f"total={timing['total_ms']:.3f}, "
|
||||
f"queue={timing['queue_ms']:.3f}, "
|
||||
f"cache={timing['cache_prepare_ms']:.3f}, "
|
||||
f"host={timing['host_prepare_ms']:.3f}, "
|
||||
f"h2d={timing['h2d_set_ms']:.3f}, "
|
||||
f"fft={timing['fft_f_ms']:.3f}, "
|
||||
f"mul={timing['mul_ms']:.3f}, "
|
||||
f"ifft={timing['ifft_ms']:.3f}, "
|
||||
f"d2h={timing['d2h_get_ms']:.3f}, "
|
||||
f"crop={timing['crop_scale_ms']:.3f}"
|
||||
)
|
||||
return p1, p2
|
||||
|
||||
@@ -0,0 +1,36 @@
|
||||
import numpy
|
||||
import timeit
|
||||
|
||||
class Timer(object):
|
||||
|
||||
def display(self):
|
||||
|
||||
tseconds = self.etime - self.stime
|
||||
|
||||
print(' [[ elapsed time:', end = "")
|
||||
if tseconds < numpy.power(10.,-6.):
|
||||
print("%.3f micro-sec." % (tseconds*numpy.power(10,6)),end = "")
|
||||
|
||||
elif tseconds < numpy.power(10.,-3.):
|
||||
print("%.3f milli-sec." % (tseconds*numpy.power(10.,3.)),end = "")
|
||||
|
||||
elif tseconds < 60.0:
|
||||
print("%.3f sec." % (tseconds),end = "")
|
||||
|
||||
else:
|
||||
m = int(tseconds/ 60.0)
|
||||
print("%d min. %.3f sec." % (m, (tseconds - 60*m)),end = "")
|
||||
|
||||
print(']]')
|
||||
|
||||
def start(self):
|
||||
self.stime = timeit.default_timer()
|
||||
|
||||
def end(self, print_time=True):
|
||||
self.etime = timeit.default_timer()
|
||||
|
||||
if print_time:
|
||||
self.display()
|
||||
|
||||
def __init__(self):
|
||||
self.start()
|
||||
@@ -19,7 +19,7 @@ from typing import Tuple
|
||||
|
||||
from .Stitcher import Stitcher
|
||||
from .tracker import Tracker
|
||||
from .HHD import nHHD
|
||||
from .hhd import nHHD
|
||||
from . import configs
|
||||
from .get_depth import Contact3D
|
||||
from .backend_availability import (check_cupy_available, check_opencl_available)
|
||||
|
||||
Reference in New Issue
Block a user