refractor HHD
This commit is contained in:
ykcheungab
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -43,10 +43,10 @@ def main():
|
||||
# 步骤 2.1：获取并拼接图像（必需）
|
||||
sensor.get_img()
|
||||
# 步骤 2.2：计算形变场（必需）；返回值表示当前帧是否检测到接触/运动
|
||||
is_contact_now = sensor.compute_deformation(check_motion=True, threshold=6)
|
||||
is_contact_now = sensor.compute_deformation(check_motion=True, threshold=0)
|
||||
# 步骤 2.3：仅在当前帧检测到接触/运动时计算接触区域
|
||||
if is_contact_now:
|
||||
sensor.compute_contact()
|
||||
# if is_contact_now:
|
||||
sensor.compute_contact()
|
||||
|
||||
|
||||
# 步骤 2.4：读取结果
|
||||
|
||||
+45
-25
@@ -27,7 +27,6 @@ INCLUDE_LAYOUT = {
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def install_nuitka():
|
||||
print("Checking Nuitka installation...")
|
||||
try:
|
||||
@@ -50,14 +49,15 @@ def get_version() -> str:
|
||||
return "0.3.1"
|
||||
|
||||
|
||||
def find_compiled_extension(output_dir: Path) -> Path:
|
||||
def find_compiled_extension(output_dir: Path, module_name: str) -> Path:
|
||||
candidates = []
|
||||
for suffix in EXTENSION_SUFFIXES:
|
||||
candidates.extend(output_dir.glob(f"*{suffix}"))
|
||||
candidates.extend(output_dir.glob(f"{module_name}*{suffix}"))
|
||||
if not candidates:
|
||||
all_files = ", ".join(p.name for p in output_dir.iterdir()) if output_dir.exists() else "(none)"
|
||||
raise RuntimeError(f"Nuitka did not produce a compiled extension. Found: {all_files}")
|
||||
# Prefer the shortest filename (usually the package extension itself)
|
||||
raise RuntimeError(
|
||||
f"Nuitka did not produce a compiled extension for {module_name}. Found: {all_files}"
|
||||
)
|
||||
return sorted(candidates, key=lambda p: len(p.name))[0]
|
||||
|
||||
|
||||
@@ -71,6 +71,40 @@ def move_wheel_to_sdk(wheel_path: Path, version: str) -> Path:
|
||||
return dest_wheel
|
||||
|
||||
|
||||
def compile_orisys(build_root: Path, env: dict[str, str]) -> Path:
|
||||
module_output_dir = build_root / "orisys"
|
||||
if module_output_dir.exists():
|
||||
shutil.rmtree(module_output_dir)
|
||||
module_output_dir.mkdir(parents=True)
|
||||
|
||||
cmd = [
|
||||
sys.executable,
|
||||
"-m",
|
||||
"nuitka",
|
||||
"--module",
|
||||
"--include-package=orisys",
|
||||
"--output-dir=" + str(module_output_dir),
|
||||
"--remove-output",
|
||||
"--no-pyi-file",
|
||||
"src/orisys",
|
||||
]
|
||||
|
||||
print(f"Running: {' '.join(cmd)}")
|
||||
subprocess.check_call(cmd, env=env)
|
||||
|
||||
ext_file = find_compiled_extension(module_output_dir, "orisys")
|
||||
print(f"Generated compiled extension for orisys: {ext_file}")
|
||||
return ext_file
|
||||
|
||||
|
||||
def copy_compiled_orisys(pkg_dir: Path, ext_file: Path) -> str:
|
||||
ext_suffix = "".join(ext_file.suffixes)
|
||||
dest_ext = pkg_dir / f"__init__{ext_suffix}"
|
||||
shutil.copy(ext_file, dest_ext)
|
||||
print(f"Copied compiled module orisys to {dest_ext}")
|
||||
return ext_suffix
|
||||
|
||||
|
||||
def copy_included_files(pkg_dir: Path) -> dict[str, int]:
|
||||
copied_counts = {}
|
||||
for group_name, group_cfg in INCLUDE_LAYOUT.items():
|
||||
@@ -100,7 +134,10 @@ def copy_included_files(pkg_dir: Path) -> dict[str, int]:
|
||||
|
||||
|
||||
def get_package_data_entries(ext_suffix: str) -> dict[str, list[str]]:
|
||||
package_data = {"orisys": [f"*{ext_suffix}"], "orisys.examples": ["*.py"]}
|
||||
package_data = {
|
||||
"orisys": [f"*{ext_suffix}"],
|
||||
"orisys.examples": ["*.py"],
|
||||
}
|
||||
for group_name, group_cfg in INCLUDE_LAYOUT.items():
|
||||
if group_name == "examples":
|
||||
continue
|
||||
@@ -127,31 +164,14 @@ def build():
|
||||
print("\n=== Compiling with Nuitka ===")
|
||||
env = os.environ.copy()
|
||||
env["PYTHONPATH"] = "src" + os.pathsep + env.get("PYTHONPATH", "")
|
||||
|
||||
cmd = [
|
||||
sys.executable, "-m", "nuitka",
|
||||
"--module",
|
||||
"--include-package=orisys",
|
||||
"--output-dir=" + str(output_dir),
|
||||
"--remove-output",
|
||||
"--no-pyi-file",
|
||||
"src/orisys",
|
||||
]
|
||||
print(f"Running: {' '.join(cmd)}")
|
||||
subprocess.check_call(cmd, env=env)
|
||||
|
||||
ext_file = find_compiled_extension(output_dir)
|
||||
ext_suffix = "".join(ext_file.suffixes)
|
||||
print(f"Generated compiled extension: {ext_file}")
|
||||
ext_file = compile_orisys(output_dir, env)
|
||||
|
||||
print("\n=== Preparing Wheel Structure ===")
|
||||
wheel_root = dist_dir / "wheel_root"
|
||||
pkg_dir = wheel_root / "orisys"
|
||||
pkg_dir.mkdir(parents=True)
|
||||
|
||||
dest_ext = pkg_dir / f"__init__{ext_suffix}"
|
||||
shutil.copy(ext_file, dest_ext)
|
||||
print(f"Copied compiled package to {dest_ext}")
|
||||
ext_suffix = copy_compiled_orisys(pkg_dir, ext_file)
|
||||
|
||||
copied_counts = copy_included_files(pkg_dir)
|
||||
print(f"Copied {copied_counts['configs']} JSON config files")
|
||||
|
||||
-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