diff --git a/data/promote/NG_1/NG_1_arrows.mp4 b/data/promote/NG_1/NG_1_arrows.mp4 deleted file mode 100644 index 42ecdf9..0000000 Binary files a/data/promote/NG_1/NG_1_arrows.mp4 and /dev/null differ diff --git a/data/promote/NG_1/NG_1_depth.mp4 b/data/promote/NG_1/NG_1_depth.mp4 deleted file mode 100644 index 138dd39..0000000 Binary files a/data/promote/NG_1/NG_1_depth.mp4 and /dev/null differ diff --git a/data/promote/NG_2/NG_2_depth.mp4 b/data/promote/NG_2/NG_2_depth.mp4 deleted file mode 100644 index 3170e23..0000000 Binary files a/data/promote/NG_2/NG_2_depth.mp4 and /dev/null differ diff --git a/data/promote/NG_3/NG_3_depth.mp4 b/data/promote/NG_3/NG_3_depth.mp4 deleted file mode 100644 index 95c12e9..0000000 Binary files a/data/promote/NG_3/NG_3_depth.mp4 and /dev/null differ diff --git a/data/promote/NGblack_Feb06_0357/NGblack_Feb06_0357.npy b/data/promote/NGblack_Feb06_0357/NGblack_Feb06_0357.npy deleted file mode 100644 index cf2d210..0000000 Binary files a/data/promote/NGblack_Feb06_0357/NGblack_Feb06_0357.npy and /dev/null differ diff --git a/data/promote/NGblack_Feb06_0357/NGblack_Feb06_0357_arrows.mp4 b/data/promote/NGblack_Feb06_0357/NGblack_Feb06_0357_arrows.mp4 deleted file mode 100644 index b82f51c..0000000 Binary files a/data/promote/NGblack_Feb06_0357/NGblack_Feb06_0357_arrows.mp4 and /dev/null differ diff --git a/data/promote/NGblack_Feb06_0357/NGblack_Feb06_0357_depth.mp4 b/data/promote/NGblack_Feb06_0357/NGblack_Feb06_0357_depth.mp4 deleted file mode 100644 index 87abf2c..0000000 Binary files a/data/promote/NGblack_Feb06_0357/NGblack_Feb06_0357_depth.mp4 and /dev/null differ diff --git a/data/promote/add_water/add_water.npy b/data/promote/add_water/add_water.npy deleted file mode 100644 index cf2d210..0000000 Binary files a/data/promote/add_water/add_water.npy and /dev/null differ diff --git a/data/promote/add_water/add_water_arrows.mp4 b/data/promote/add_water/add_water_arrows.mp4 deleted file mode 100644 index 9b889f0..0000000 Binary files a/data/promote/add_water/add_water_arrows.mp4 and /dev/null differ diff --git a/data/promote/add_water/add_water_depth.mp4 b/data/promote/add_water/add_water_depth.mp4 deleted file mode 100644 index 050e27a..0000000 Binary files a/data/promote/add_water/add_water_depth.mp4 and /dev/null differ diff --git a/data/promote/egg_big/egg_big_arrows.mp4 b/data/promote/egg_big/egg_big_arrows.mp4 deleted file mode 100644 index 8e07ed8..0000000 Binary files a/data/promote/egg_big/egg_big_arrows.mp4 and /dev/null differ diff --git a/data/promote/egg_big/egg_big_depth.mp4 b/data/promote/egg_big/egg_big_depth.mp4 deleted file mode 100644 index ef4bc0c..0000000 Binary files a/data/promote/egg_big/egg_big_depth.mp4 and /dev/null differ diff --git a/data/promote/egg_plate_1shot/egg_plate_1shot_arrows.mp4 b/data/promote/egg_plate_1shot/egg_plate_1shot_arrows.mp4 deleted file mode 100644 index 47e466d..0000000 Binary files a/data/promote/egg_plate_1shot/egg_plate_1shot_arrows.mp4 and /dev/null differ diff --git a/data/promote/egg_plate_1shot/egg_plate_1shot_depth.mp4 b/data/promote/egg_plate_1shot/egg_plate_1shot_depth.mp4 deleted file mode 100644 index 9c91af5..0000000 Binary files a/data/promote/egg_plate_1shot/egg_plate_1shot_depth.mp4 and /dev/null differ diff --git a/data/promote/js2/js2_arrows.mp4 b/data/promote/js2/js2_arrows.mp4 deleted file mode 100644 index 0bb96dd..0000000 Binary files a/data/promote/js2/js2_arrows.mp4 and /dev/null differ diff --git a/data/promote/js2/js2_depth.mp4 b/data/promote/js2/js2_depth.mp4 deleted file mode 100644 index 3728e55..0000000 Binary files a/data/promote/js2/js2_depth.mp4 and /dev/null differ diff --git a/data/promote/pin1/pin1_depth.mp4 b/data/promote/pin1/pin1_depth.mp4 deleted file mode 100644 index bf70a6a..0000000 Binary files a/data/promote/pin1/pin1_depth.mp4 and /dev/null differ diff --git a/data/promote/pin2/pin2_depth.mp4 b/data/promote/pin2/pin2_depth.mp4 deleted file mode 100644 index 1d7a93e..0000000 Binary files a/data/promote/pin2/pin2_depth.mp4 and /dev/null differ diff --git a/data/promote/pin3/pin3.npy b/data/promote/pin3/pin3.npy deleted file mode 100644 index 068be20..0000000 Binary files a/data/promote/pin3/pin3.npy and /dev/null differ diff --git a/data/promote/pin3/pin3_arrows.mp4 b/data/promote/pin3/pin3_arrows.mp4 deleted file mode 100644 index 96a98e1..0000000 Binary files a/data/promote/pin3/pin3_arrows.mp4 and /dev/null differ diff --git a/data/promote/pin3/pin3_depth.mp4 b/data/promote/pin3/pin3_depth.mp4 deleted file mode 100644 index 3c915bc..0000000 Binary files a/data/promote/pin3/pin3_depth.mp4 and /dev/null differ diff --git a/data/promote/pin3/pin3_pyqt.mp4 b/data/promote/pin3/pin3_pyqt.mp4 deleted file mode 100644 index 9647432..0000000 Binary files a/data/promote/pin3/pin3_pyqt.mp4 and /dev/null differ diff --git a/data/promote/pin_l2_1/pin_l2_1.npy b/data/promote/pin_l2_1/pin_l2_1.npy deleted file mode 100644 index 068be20..0000000 Binary files a/data/promote/pin_l2_1/pin_l2_1.npy and /dev/null differ diff --git a/data/promote/pin_l2_1/pin_l2_1_arrows.mp4 b/data/promote/pin_l2_1/pin_l2_1_arrows.mp4 deleted file mode 100644 index 96a98e1..0000000 Binary files a/data/promote/pin_l2_1/pin_l2_1_arrows.mp4 and /dev/null differ diff --git a/data/promote/pin_l2_1/pin_l2_1_depth.mp4 b/data/promote/pin_l2_1/pin_l2_1_depth.mp4 deleted file mode 100644 index 3c915bc..0000000 Binary files a/data/promote/pin_l2_1/pin_l2_1_depth.mp4 and /dev/null differ diff --git a/data/promote/pin_l2_1/pin_l2_1_pyqt.mp4 b/data/promote/pin_l2_1/pin_l2_1_pyqt.mp4 deleted file mode 100644 index bdeb370..0000000 Binary files a/data/promote/pin_l2_1/pin_l2_1_pyqt.mp4 and /dev/null differ diff --git a/data/promote/pin_l2_1_arrows.mp4 b/data/promote/pin_l2_1_arrows.mp4 deleted file mode 100644 index 566a00f..0000000 Binary files a/data/promote/pin_l2_1_arrows.mp4 and /dev/null differ diff --git a/data/promote/pin_l2_1_depth.mp4 b/data/promote/pin_l2_1_depth.mp4 deleted file mode 100644 index 566a00f..0000000 Binary files a/data/promote/pin_l2_1_depth.mp4 and /dev/null differ diff --git a/data/promote/pin_l2_1_pyqt.mp4 b/data/promote/pin_l2_1_pyqt.mp4 deleted file mode 100644 index 566a00f..0000000 Binary files a/data/promote/pin_l2_1_pyqt.mp4 and /dev/null differ diff --git a/data/promote/pin_l2_2/pin_l2_2_arrows.mp4 b/data/promote/pin_l2_2/pin_l2_2_arrows.mp4 deleted file mode 100644 index b3f74a4..0000000 Binary files a/data/promote/pin_l2_2/pin_l2_2_arrows.mp4 and /dev/null differ diff --git a/data/promote/pin_l2_2/pin_l2_2_depth.mp4 b/data/promote/pin_l2_2/pin_l2_2_depth.mp4 deleted file mode 100644 index 3f8eb68..0000000 Binary files a/data/promote/pin_l2_2/pin_l2_2_depth.mp4 and /dev/null differ diff --git a/data/promote/pulse_1/pulse_1_arrows.mp4 b/data/promote/pulse_1/pulse_1_arrows.mp4 deleted file mode 100644 index 2762c00..0000000 Binary files a/data/promote/pulse_1/pulse_1_arrows.mp4 and /dev/null differ diff --git a/data/promote/pulse_1/pulse_1_depth.mp4 b/data/promote/pulse_1/pulse_1_depth.mp4 deleted file mode 100644 index a48dd4b..0000000 Binary files a/data/promote/pulse_1/pulse_1_depth.mp4 and /dev/null differ diff --git a/data/promote/pulse_test_2_depth.mp4 b/data/promote/pulse_test_2_depth.mp4 deleted file mode 100644 index 29a0816..0000000 Binary files a/data/promote/pulse_test_2_depth.mp4 and /dev/null differ diff --git a/data/promote/rotor1/rotor1_arrows.mp4 b/data/promote/rotor1/rotor1_arrows.mp4 deleted file mode 100644 index 525c367..0000000 Binary files a/data/promote/rotor1/rotor1_arrows.mp4 and /dev/null differ diff --git a/data/promote/rotor1/rotor1_depth.mp4 b/data/promote/rotor1/rotor1_depth.mp4 deleted file mode 100644 index 916e31d..0000000 Binary files a/data/promote/rotor1/rotor1_depth.mp4 and /dev/null differ diff --git a/data/promote/rotor2_Feb05_0834/rotor2_Feb05_0834.npy b/data/promote/rotor2_Feb05_0834/rotor2_Feb05_0834.npy deleted file mode 100644 index cf2d210..0000000 Binary files a/data/promote/rotor2_Feb05_0834/rotor2_Feb05_0834.npy and /dev/null differ diff --git a/data/promote/rotor2_Feb05_0834/rotor2_Feb05_0834_arrows.mp4 b/data/promote/rotor2_Feb05_0834/rotor2_Feb05_0834_arrows.mp4 deleted file mode 100644 index 9f9356e..0000000 Binary files a/data/promote/rotor2_Feb05_0834/rotor2_Feb05_0834_arrows.mp4 and /dev/null differ diff --git a/data/promote/rotor2_Feb05_0834/rotor2_Feb05_0834_depth.mp4 b/data/promote/rotor2_Feb05_0834/rotor2_Feb05_0834_depth.mp4 deleted file mode 100644 index 9f9356e..0000000 Binary files a/data/promote/rotor2_Feb05_0834/rotor2_Feb05_0834_depth.mp4 and /dev/null differ diff --git a/data/promote/rotor2_Feb05_0834/rotor2_Feb05_0834_pyqt.mp4 b/data/promote/rotor2_Feb05_0834/rotor2_Feb05_0834_pyqt.mp4 deleted file mode 100644 index 9f9356e..0000000 Binary files a/data/promote/rotor2_Feb05_0834/rotor2_Feb05_0834_pyqt.mp4 and /dev/null differ diff --git a/data/promote/rotor3_Feb05_0836/rotor3_Feb05_0836_arrows.mp4 b/data/promote/rotor3_Feb05_0836/rotor3_Feb05_0836_arrows.mp4 deleted file mode 100644 index 09ad013..0000000 Binary files a/data/promote/rotor3_Feb05_0836/rotor3_Feb05_0836_arrows.mp4 and /dev/null differ diff --git a/data/promote/rotor3_Feb05_0836/rotor3_Feb05_0836_depth.mp4 b/data/promote/rotor3_Feb05_0836/rotor3_Feb05_0836_depth.mp4 deleted file mode 100644 index b96cf11..0000000 Binary files a/data/promote/rotor3_Feb05_0836/rotor3_Feb05_0836_depth.mp4 and /dev/null differ diff --git a/data/promote/test_Feb1_depth.mp4 b/data/promote/test_Feb1_depth.mp4 deleted file mode 100644 index c862236..0000000 Binary files a/data/promote/test_Feb1_depth.mp4 and /dev/null differ diff --git a/data/promote/test_pulse/test_pulse_arrows.mp4 b/data/promote/test_pulse/test_pulse_arrows.mp4 deleted file mode 100644 index eed191d..0000000 Binary files a/data/promote/test_pulse/test_pulse_arrows.mp4 and /dev/null differ diff --git a/data/promote/test_pulse/test_pulse_depth.mp4 b/data/promote/test_pulse/test_pulse_depth.mp4 deleted file mode 100644 index 706cffe..0000000 Binary files a/data/promote/test_pulse/test_pulse_depth.mp4 and /dev/null differ diff --git a/data/promote/unsorted/WIN_20260206_15_24_40_Pro.mp4 b/data/promote/unsorted/WIN_20260206_15_24_40_Pro.mp4 deleted file mode 100644 index 361e69f..0000000 Binary files a/data/promote/unsorted/WIN_20260206_15_24_40_Pro.mp4 and /dev/null differ diff --git a/examples/basic_pipeline.py b/examples/basic_pipeline.py index 0cf40b2..7404637 100644 --- a/examples/basic_pipeline.py +++ b/examples/basic_pipeline.py @@ -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：读取结果 diff --git a/scripts/build_nuitka.py b/scripts/build_nuitka.py index 0e4d286..99b82b2 100644 --- a/scripts/build_nuitka.py +++ b/scripts/build_nuitka.py @@ -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") diff --git a/src/orisys/HHD.py b/src/orisys/HHD.py deleted file mode 100644 index 3c6f919..0000000 --- a/src/orisys/HHD.py +++ /dev/null @@ -1,2072 +0,0 @@ -''' -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 signal, 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.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 - - # 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) - - # 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 - # 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) - - # 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() - - - -# --------------------------------------------------------------------------- - -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) - # 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 - - # def solve_opencl(self, f, verbose=False, profile=True): - # """ - # 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) - - diff --git a/src/orisys/hhd/__init__.py b/src/orisys/hhd/__init__.py new file mode 100644 index 0000000..5a2f11c --- /dev/null +++ b/src/orisys/hhd/__init__.py @@ -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"] diff --git a/src/orisys/hhd/api.py b/src/orisys/hhd/api.py new file mode 100644 index 0000000..e03c2c4 --- /dev/null +++ b/src/orisys/hhd/api.py @@ -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() diff --git a/src/orisys/hhd/grid.py b/src/orisys/hhd/grid.py new file mode 100644 index 0000000..cdcf20b --- /dev/null +++ b/src/orisys/hhd/grid.py @@ -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() + + # -------------------------------------------------------------------------- + # -------------------------------------------------------------------------- diff --git a/src/orisys/hhd/kernels.py b/src/orisys/hhd/kernels.py new file mode 100644 index 0000000..7261f69 --- /dev/null +++ b/src/orisys/hhd/kernels.py @@ -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__ = [] diff --git a/src/orisys/hhd/poisson.py b/src/orisys/hhd/poisson.py new file mode 100644 index 0000000..a1a6a17 --- /dev/null +++ b/src/orisys/hhd/poisson.py @@ -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 + \ No newline at end of file diff --git a/src/orisys/hhd/timing.py b/src/orisys/hhd/timing.py new file mode 100644 index 0000000..9253c5a --- /dev/null +++ b/src/orisys/hhd/timing.py @@ -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() \ No newline at end of file diff --git a/src/orisys/sensor.py b/src/orisys/sensor.py index f822787..461e745 100644 --- a/src/orisys/sensor.py +++ b/src/orisys/sensor.py @@ -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)