bluepulsar
2013 lines
79 KiB
Python
2013 lines
79 KiB
Python
'''
|
|
Copyright (c) 2015, Harsh Bhatia (bhatia4@llnl.gov)
|
|
All rights reserved.
|
|
|
|
Redistribution and use in source and binary forms, with or without
|
|
modification, are permitted provided that the following conditions are met:
|
|
1. Redistributions of source code must retain the above copyright notice, this
|
|
list of conditions and the following disclaimer.
|
|
2. Redistributions in binary form must reproduce the above copyright notice,
|
|
this list of conditions and the following disclaimer in the documentation
|
|
and/or other materials provided with the distribution.
|
|
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
|
|
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
|
|
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
|
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
|
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
|
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
|
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
|
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
|
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
|
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
|
'''
|
|
|
|
import numpy
|
|
import timeit
|
|
import sys
|
|
import cv2
|
|
from scipy import spatial
|
|
|
|
from .backend_availability import (
|
|
CUPY_AVAILABLE, OPENCL_AVAILABLE, KERNEL_AVAILABLE,
|
|
cp, cp_signal, _torque_kernel,
|
|
cl, cla, cl_rfftn, cl_irfftn
|
|
)
|
|
# ---------------------------------------------------------------------------
|
|
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.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':
|
|
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':
|
|
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=False)
|
|
|
|
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 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=False)
|
|
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()))
|
|
|
|
with self.stream1:
|
|
nD_gpu = self.psolver._solve_gpu_internal(div_gpu)
|
|
with self.stream2:
|
|
nRu_gpu = self.psolver._solve_gpu_internal(curlw_gpu)
|
|
self.stream1.synchronize()
|
|
self.stream2.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()))
|
|
|
|
# Store GPU arrays
|
|
self._d_gpu = d_gpu
|
|
self._h_gpu = h_gpu
|
|
self._r_gpu = r_gpu
|
|
self._nRu_gpu = nRu_gpu
|
|
|
|
# Also store CPU copies for compatibility
|
|
# self.d = cp.asnumpy(d_gpu)
|
|
# self.h = cp.asnumpy(h_gpu)
|
|
# self.r = cp.asnumpy(r_gpu)
|
|
# self.nRu = cp.asnumpy(nRu_gpu)
|
|
# self.nD = cp.asnumpy(nD_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)
|
|
|
|
self.Sn = numpy.sum(numpy.absolute(self.d))
|
|
# print("Sn is:", self.Sn)
|
|
|
|
# self.Stx = abs(numpy.sum(self.h[:, :, 0]))
|
|
# self.Sty = abs(numpy.sum(self.h[:, :, 1]))
|
|
self.Stx = numpy.sum(self.h[:, :, 0])
|
|
self.Sty = numpy.sum(self.h[:, :, 1])
|
|
# print("Stx is:", self.Stx)
|
|
# print("Sty is:", self.Sty)
|
|
|
|
nR_max_loc = numpy.unravel_index(numpy.argmax(self.nRu, axis = None), self.nRu.shape)
|
|
nR_min_loc = numpy.unravel_index(numpy.argmin(self.nRu, axis = None), self.nRu.shape)
|
|
|
|
# print(nR_max_loc)
|
|
|
|
r_p = grid - nR_max_loc
|
|
r_n = grid - nR_min_loc
|
|
|
|
self.Stor = numpy.sum(numpy.cross(r_p, self.r)) + numpy.sum(numpy.cross(r_n, self.r))
|
|
# print("Stor is:", self.Stor)
|
|
|
|
output = numpy.array([self.Sn, self.Stx, self.Sty, self.Stor])
|
|
|
|
flag = False
|
|
if self.Sn >=300 and self.Stx >= 300 and self.Sty >= 300 and numpy.abs(self.Stor) >= 1000:
|
|
flag = True
|
|
# print(flag)
|
|
return flag,output
|
|
# return self.d, self.h, numpy.cross(r_p, self.r) , numpy.cross(r_n, self.r)
|
|
|
|
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)
|
|
|
|
# Keep data on GPU - avoid transfers by computing directly
|
|
# d_gpu = cp.asarray(self.d, dtype=cp.float32)
|
|
# h_gpu = cp.asarray(self.h, dtype=cp.float32)
|
|
# r_gpu = cp.asarray(self.r, dtype=cp.float32)
|
|
# nRu_gpu = cp.asarray(self.nRu, dtype=cp.float32)
|
|
|
|
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]
|
|
|
|
# Compute torque with fused sum
|
|
# 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
|
|
# ))
|
|
# Replace lines 413-416:
|
|
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))
|
|
|
|
self.Sn, self.Stx, self.Sty, self.Stor = Sn, Stx, Sty, Stor
|
|
output = numpy.array([Sn, Stx, Sty, Stor], dtype=numpy.float32)
|
|
|
|
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()
|
|
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
|
|
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()
|
|
|
|
# ---------------------------------------------------------------------------
|
|
|
|
# Try to import CuPy for GPU acceleration
|
|
# try:
|
|
# import cupy as cp
|
|
# CUPY_AVAILABLE = True
|
|
# except ImportError:
|
|
# cp = None
|
|
# CUPY_AVAILABLE = False
|
|
|
|
# ------------------------------------------------------------------------------
|
|
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()
|
|
|
|
# --------------------------------------------------------------------------
|
|
# --------------------------------------------------------------------------
|
|
|
|
|
|
# import numpy
|
|
# from scipy import signal, spatial
|
|
# # from .timer import Timer
|
|
|
|
# # Try to import CuPy for GPU acceleration
|
|
# try:
|
|
# import cupy as cp
|
|
# import cupyx.scipy.signal as cp_signal
|
|
# CUPY_AVAILABLE = True
|
|
# except ImportError:
|
|
# cp = None
|
|
# cp_signal = None
|
|
# CUPY_AVAILABLE = False
|
|
|
|
# Try to import OpenCL + pyvkfft for FFT-based convolution
|
|
|
|
|
|
|
|
'''
|
|
# ------------------------------------------------------------------------------
|
|
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)
|
|
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."""
|
|
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
|
|
) {
|
|
int i = get_global_id(0);
|
|
int b = get_global_id(1);
|
|
if (i < n) {
|
|
float2 gv = g[i];
|
|
int idx = b * n + i;
|
|
float2 a = fft_batch[idx];
|
|
fft_batch[idx] = (float2)(a.x * gv.x - a.y * gv.y,
|
|
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
|
|
|
|
@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()
|
|
f1_host = numpy.ascontiguousarray(f1, dtype=numpy.float32)
|
|
f2_host = numpy.ascontiguousarray(f2, dtype=numpy.float32)
|
|
full_shape = tuple(2 * s - 1 for s in f1_host.shape)
|
|
timing['host_prepare_ms'] = (timeit.default_timer() - t0) * 1000.0
|
|
|
|
cache_valid = (
|
|
self._opencl_cache is not None and
|
|
self._opencl_cache['fshape'] == f1_host.shape 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 = [f1_host.shape[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 = f1_host.shape[0] * f1_host.shape[1]
|
|
dst_plane = full_shape[0] * full_shape[1]
|
|
|
|
self._opencl_cache = {
|
|
'fshape': f1_host.shape,
|
|
'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,
|
|
# Stacked (batch=2) buffers for true batched FFT
|
|
'f_batch_dev': cla.zeros(queue, (2,) + full_shape, dtype=numpy.float32),
|
|
'f_stage_batch_dev': cla.empty(queue, (2,) + f1_host.shape, 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,) + f1_host.shape, dtype=numpy.float32),
|
|
'p_batch_host': numpy.empty((2,) + f1_host.shape, dtype=numpy.float32),
|
|
'f_stage_batch_host': numpy.empty((2,) + f1_host.shape, 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: single transfer for both inputs, then batch scatter into padded buffer
|
|
t0 = timeit.default_timer()
|
|
f_stage_host = c['f_stage_batch_host']
|
|
f_stage_host[0] = f1_host
|
|
f_stage_host[1] = f2_host
|
|
c['f_stage_batch_dev'].set(f_stage_host, queue=queue)
|
|
scatter_batch = self._get_opencl_scatter_batch_kernel(queue)
|
|
scatter_batch(
|
|
queue, (fshape_c[1], fshape_c[0], 2), None,
|
|
c['f_stage_batch_dev'].data,
|
|
c['f_batch_dev'].data,
|
|
numpy.int32(fshape_c[1]),
|
|
numpy.int32(fshape_c[0]),
|
|
numpy.int32(full_shape_c[1]),
|
|
numpy.int32(c['src_plane']),
|
|
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
|
|
t0 = timeit.default_timer()
|
|
batch_mul = self._get_opencl_batch_mul_kernel(queue)
|
|
batch_mul(
|
|
queue, (c['fft_n_elements'], 2), None,
|
|
fft_batch.data,
|
|
c['g_fft_dev'].data,
|
|
numpy.int32(c['fft_n_elements'])
|
|
)
|
|
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: single transfer for both results
|
|
t0 = timeit.default_timer()
|
|
cl.enqueue_copy(queue, c['p_batch_host'], c['p_batch_dev'].data, is_blocking=True)
|
|
p1 = c['p_batch_host'][0]
|
|
p2 = c['p_batch_host'][1]
|
|
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
|
|
|
|
# def solve_opencl(self, f, verbose=False, profile=False):
|
|
# """
|
|
# OpenCL Poisson solver using pyvkfft FFT convolution.
|
|
# Replaces scipy.signal.fftconvolve(..., mode='same') for structured grids.
|
|
# 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(f, verbose)
|
|
|
|
# if not self.ready:
|
|
# self.prepare(verbose)
|
|
|
|
# fshape = f.shape
|
|
# gtimer = Timer()
|
|
# if verbose:
|
|
# print('\nSolving Poisson Eq. with OpenCL FFT')
|
|
|
|
# if self.ptype == 'G':
|
|
# if self.gdims != fshape:
|
|
# print(self.gdims, fshape)
|
|
# raise ValueError("Shape of function should match shape of grid")
|
|
|
|
# if self.stype == 'F':
|
|
# 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()
|
|
# f_host = numpy.ascontiguousarray(f, dtype=numpy.float32)
|
|
# full_shape = tuple(2 * s - 1 for s in f_host.shape)
|
|
# timing['host_prepare_ms'] = (timeit.default_timer() - t0) * 1000.0
|
|
|
|
# cache_valid = (
|
|
# self._opencl_cache is not None and
|
|
# self._opencl_cache['fshape'] == f_host.shape and
|
|
# self._opencl_cache['full_shape'] == full_shape and
|
|
# 'f_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 = [f_host.shape[i] - 1 for i in range(self.dim)]
|
|
# same_slices = tuple(
|
|
# slice(same_starts[i], same_starts[i] + f_host.shape[i])
|
|
# for i in range(self.dim)
|
|
# )
|
|
# fft_shape = tuple(list(full_shape[:-1]) + [full_shape[-1] // 2 + 1])
|
|
# self._opencl_cache = {
|
|
# 'fshape': f_host.shape,
|
|
# 'full_shape': full_shape,
|
|
# 'fft_shape': fft_shape,
|
|
# 'g_fft_dev': g_fft_dev,
|
|
# 'same_slices': same_slices,
|
|
# 'same_starts': same_starts,
|
|
# 'gdx_prod': numpy.prod(self.gdx),
|
|
# # Zero-init so padding stays zero across frames
|
|
# 'f_dev': cla.zeros(queue, full_shape, dtype=numpy.float32),
|
|
# # Small staging buffer for input-sized H2D
|
|
# 'f_stage_dev': cla.empty(queue, f_host.shape, dtype=numpy.float32),
|
|
# 'f_fft_dev': cla.empty(queue, fft_shape, dtype=numpy.complex64),
|
|
# 'conv_full_dev': cla.empty(queue, full_shape, dtype=numpy.float32),
|
|
# 'p_dev': cla.empty(queue, f_host.shape, dtype=numpy.float32),
|
|
# 'p_host': numpy.empty(f_host.shape, dtype=numpy.float32),
|
|
# }
|
|
# timing['cache_prepare_ms'] = (timeit.default_timer() - t0) * 1000.0
|
|
|
|
# # H2D: upload only input-sized data, scatter into padded device buffer
|
|
# t0 = timeit.default_timer()
|
|
# self._opencl_cache['f_stage_dev'].set(f_host, queue=queue)
|
|
# if self.dim == 2:
|
|
# scatter_kern = self._get_opencl_scatter_kernel(queue)
|
|
# fshape_c = self._opencl_cache['fshape']
|
|
# full_shape_c = self._opencl_cache['full_shape']
|
|
# scatter_kern(
|
|
# queue, (fshape_c[1], fshape_c[0]), None,
|
|
# self._opencl_cache['f_stage_dev'].data,
|
|
# self._opencl_cache['f_dev'].data,
|
|
# numpy.int32(fshape_c[1]),
|
|
# numpy.int32(fshape_c[0]),
|
|
# numpy.int32(full_shape_c[1])
|
|
# )
|
|
# else:
|
|
# # 3D fallback: host-side pad + full upload
|
|
# f_pad = numpy.zeros(full_shape, dtype=numpy.float32)
|
|
# f_slices = tuple(slice(0, s) for s in f_host.shape)
|
|
# f_pad[f_slices] = f_host
|
|
# self._opencl_cache['f_dev'].set(f_pad, queue=queue)
|
|
# if profile:
|
|
# queue.finish()
|
|
# timing['h2d_set_ms'] = (timeit.default_timer() - t0) * 1000.0
|
|
|
|
# t0 = timeit.default_timer()
|
|
# f_fft_dev = cl_rfftn(self._opencl_cache['f_dev'], self._opencl_cache['f_fft_dev'])
|
|
# if profile:
|
|
# queue.finish()
|
|
# timing['fft_f_ms'] = (timeit.default_timer() - t0) * 1000.0
|
|
|
|
# t0 = timeit.default_timer()
|
|
# f_fft_dev *= self._opencl_cache['g_fft_dev']
|
|
# if profile:
|
|
# queue.finish()
|
|
# timing['mul_ms'] = (timeit.default_timer() - t0) * 1000.0
|
|
|
|
# t0 = timeit.default_timer()
|
|
# conv_full_dev = cl_irfftn(f_fft_dev, self._opencl_cache['conv_full_dev'])
|
|
# if profile:
|
|
# queue.finish()
|
|
# timing['ifft_ms'] = (timeit.default_timer() - t0) * 1000.0
|
|
|
|
# # Optimized path: use OpenCL kernel to crop and scale on device, then transfer only cropped result.
|
|
# t0 = timeit.default_timer()
|
|
# p_dev = self._opencl_cache['p_dev']
|
|
# same_starts = self._opencl_cache['same_starts']
|
|
# fshape = self._opencl_cache['fshape']
|
|
# full_shape_cache = self._opencl_cache['full_shape']
|
|
|
|
# if self.dim == 2:
|
|
# # Use custom kernel for 2D cropping and scaling
|
|
# crop_kernel = self._get_opencl_crop_kernel(queue)
|
|
# crop_kernel(
|
|
# queue, (fshape[1], fshape[0]), None,
|
|
# conv_full_dev.data,
|
|
# p_dev.data,
|
|
# numpy.int32(full_shape_cache[1]), # src_width
|
|
# numpy.int32(fshape[1]), # dst_width
|
|
# numpy.int32(fshape[0]), # dst_height
|
|
# numpy.int32(same_starts[1]), # start_x
|
|
# numpy.int32(same_starts[0]), # start_y
|
|
# numpy.float32(self._opencl_cache['gdx_prod']) # scale
|
|
# )
|
|
# if profile:
|
|
# queue.finish()
|
|
# timing['crop_scale_ms'] = (timeit.default_timer() - t0) * 1000.0
|
|
|
|
# t0 = timeit.default_timer()
|
|
# p = p_dev.get(ary=self._opencl_cache['p_host'])
|
|
# timing['d2h_get_ms'] = (timeit.default_timer() - t0) * 1000.0
|
|
# else:
|
|
# # Fallback for 3D: transfer full, crop on CPU
|
|
# conv_full_dev *= self._opencl_cache['gdx_prod']
|
|
# queue.finish()
|
|
# timing['crop_scale_ms'] = (timeit.default_timer() - t0) * 1000.0
|
|
|
|
# t0 = timeit.default_timer()
|
|
# conv_full = conv_full_dev.get()
|
|
# timing['d2h_get_ms'] = (timeit.default_timer() - t0) * 1000.0
|
|
|
|
# t0 = timeit.default_timer()
|
|
# p = conv_full[self._opencl_cache['same_slices']].copy()
|
|
# timing['crop_scale_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) finished')
|
|
# if profile:
|
|
# print(
|
|
# "OpenCL 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}"
|
|
# )
|
|
# gtimer.end()
|
|
# return p
|
|
|
|
# elif self.stype == 'S':
|
|
# raise SyntaxError('Spatial solver not suggested for structured grids')
|
|
|
|
# elif self.ptype == 'P':
|
|
# if verbose:
|
|
# print('OpenCL solver not optimized for unstructured points, using CPU')
|
|
# return self.solve(f, verbose)
|
|
|
|
|