orisys-tech/orisys-sdk
3
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

combined HHD

Browse Source
This commit is contained in:
bluepulsar
2026-01-04 20:14:10 +08:00
parent 3d805c69db
commit 9b2545b4ee
4 changed files with 1128 additions and 270 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+1
View File
@@ -6,3 +6,4 @@ __pycache__/
orisys.egg-info/ orisys.egg-info/
/.cursor/ /.cursor/
/dist_pyd/ /dist_pyd/
/dist_single_pyd/
src/orisys/HHD.py
+857
View File
@@ -0,0 +1,857 @@
'''
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
from .nPoisson import PoissonSolver
from .StructuredGrid import StructuredGrid
import timeit
import sys
# from .UnstructuredGrid import UnstructuredGrid
# import matplotlib.pyplot as plt
# Try to import CuPy for GPU acceleration
try:
import cupy as cp
CUPY_AVAILABLE = True
# Try to import CuPy streams for parallel execution
try:
from cupy.cuda import Stream
STREAMS_AVAILABLE = True
except ImportError:
Stream = None
STREAMS_AVAILABLE = False
# Create fused elementwise kernel for torque computation
# Computes cross products and combines them in a single kernel pass
try:
_torque_kernel = cp.ElementwiseKernel(
'float32 r_x, float32 r_y, float32 grid_0, float32 grid_1, '
'float32 max_0, float32 max_1, float32 min_0, float32 min_1',
'float32 torque',
'''
float cross_p = (grid_0 - max_0) * r_y - (grid_1 - max_1) * r_x;
float cross_n = (grid_0 - min_0) * r_y - (grid_1 - min_1) * r_x;
torque = cross_p + cross_n;
''',
'torque_kernel'
)
KERNEL_AVAILABLE = True
except Exception:
_torque_kernel = None
KERNEL_AVAILABLE = False
except ImportError:
cp = None
CUPY_AVAILABLE = False
STREAMS_AVAILABLE = False
KERNEL_AVAILABLE = False
Stream = None
_torque_kernel = None
# ---------------------------------------------------------------------------
class nHHD(object):
# constructor
def __init__(self, **kwargs):
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)
# elif 'points' in args:
# if 'simplices' in args:
# self.mesh = UnstructuredGrid(vertices=kwargs['points'], simplices=kwargs['simplices'])
# else:
# self.mesh = UnstructuredGrid(vertices=kwargs['points'])
# self.psolver = PoissonSolver(points=self.mesh.vertices, pvolumes=self.mesh.pvolumes)
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.psolver.prepare(False)
# create the 3-component decomposition
def decompose(self, vfield, verbose=0):
if vfield.shape[-1] != self.dim:
raise ValueError("nHHD.decompose requires a valid-dimensional vector field")
if verbose > 0:
print( '\n ==> nhhd.decompose', vfield.shape)
if self.dim == 2:
# compute div curl
(self.div, self.curlw) = self.mesh.divcurl(vfield)
if verbose > 0:
print( 'div =', self.div.shape, self.div.min(), self.div.max())
print( 'curlw =', self.curlw.shape, self.curlw.min(), self.curlw.max())
# compute potentials
self.nD = self.psolver.solve_gpu(self.div, verbose > 0) #1)
if verbose > 0:
print( 'nD =', self.nD.shape, self.nD.min(), self.nD.max())
self.nRu = self.psolver.solve_gpu(self.curlw, verbose > 0) #1)
if verbose > 0:
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.psolver.solve_gpu(self.div, verbose > 0) #1)
if verbose > 0:
print( 'nD =', self.nD.shape, self.nD.min(), self.nD.max())
self.nRu = self.psolver.solve_gpu(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.psolver.solve_gpu(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.psolver.solve_gpu(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_gpu(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_gpu', 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 = dudx + dvdy
curlw_gpu = 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()))
# Compute potentials on GPU (solve_gpu returns numpy, so we convert back)
nD_cpu = self.psolver.solve_gpu(cp.asnumpy(div_gpu), verbose > 0)
nRu_cpu = self.psolver.solve_gpu(cp.asnumpy(curlw_gpu), verbose > 0)
# Keep potentials on GPU
nD_gpu = cp.asarray(nD_cpu, dtype=cp.float32)
nRu_gpu = cp.asarray(nRu_cpu, dtype=cp.float32)
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 *= -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 = d_gpu + r_gpu
h_gpu = vfield_gpu - h_gpu
if verbose > 0:
print('h =', h_gpu.shape, float(h_gpu.min()), float(h_gpu.max()))
# Store GPU arrays
self._d_gpu = d_gpu
self._h_gpu = h_gpu
self._r_gpu = r_gpu
self._nRu_gpu = nRu_gpu
# Also store CPU copies for compatibility
self.d = cp.asnumpy(d_gpu)
self.h = cp.asnumpy(h_gpu)
self.r = cp.asnumpy(r_gpu)
self.nRu = cp.asnumpy(nRu_gpu)
self.nD = nD_cpu
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(self, vfield, verbose=False):
mgrid = numpy.mgrid[0: vfield.shape[0], 0: vfield.shape[1]]
grid = numpy.stack((mgrid[0], mgrid[1]), axis =2)
# print(grid.shape)
self.decompose(vfield, verbose)
self.Sn = numpy.sum(numpy.absolute(self.d))
# print("Sn is:", self.Sn)
# self.Stx = abs(numpy.sum(self.h[:, :, 0]))
# self.Sty = abs(numpy.sum(self.h[:, :, 1]))
self.Stx = numpy.sum(self.h[:, :, 0])
self.Sty = numpy.sum(self.h[:, :, 1])
# print("Stx is:", self.Stx)
# print("Sty is:", self.Sty)
nR_max_loc = numpy.unravel_index(numpy.argmax(self.nRu, axis = None), self.nRu.shape)
nR_min_loc = numpy.unravel_index(numpy.argmin(self.nRu, axis = None), self.nRu.shape)
# print(nR_max_loc)
r_p = grid - nR_max_loc
r_n = grid - nR_min_loc
self.Stor = numpy.sum(numpy.cross(r_p, self.r)) + numpy.sum(numpy.cross(r_n, self.r))
# print("Stor is:", self.Stor)
output = numpy.array([self.Sn, self.Stx, self.Sty, self.Stor])
flag = False
if self.Sn >=300 and self.Stx >= 300 and self.Sty >= 300 and numpy.abs(self.Stor) >= 1000:
flag = True
# print(flag)
return flag,output
# return self.d, self.h, numpy.cross(r_p, self.r) , numpy.cross(r_n, self.r)
def _init_force_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_gpu_v2(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(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)
# Use reduction kernels for better performance
# Compute Sn using L1 norm (sum of absolute values)
Sn = float(cp.linalg.norm(d_gpu.ravel(), ord=1))
# 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
))
self.Sn, self.Stx, self.Sty, self.Stor = Sn, Stx, Sty, Stor
output = numpy.array([Sn, Stx, Sty, Stor], dtype=numpy.float32)
flag = (Sn >= 300 and Stx >= 300 and Sty >= 300 and abs(Stor) >= 1000)
return flag, output
def clear_force_gpu_cache(self):
"""Clear the GPU cache for force computation to free memory."""
if hasattr(self, '_force_gpu_cache'):
self._force_gpu_cache = None
if CUPY_AVAILABLE:
cp.get_default_memory_pool().free_all_blocks()
def get_shear(self):
return self.h
def get_normal(self):
return self.d
# ---------------------------------------------------------------------------
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()
# --------------------------------------------------------------------------
# --------------------------------------------------------------------------
src/orisys/nHHD.py
+268 -268
View File
@@ -288,87 +288,87 @@ class nHHD(object):
self.decompose(vfield, verbose) self.decompose(vfield, verbose)
return None return None
def decompose_batch(self, vfield, verbose=0): # def decompose_batch(self, vfield, verbose=0):
if vfield.shape[-1] != self.dim: # if vfield.shape[-1] != self.dim:
raise ValueError("nHHD.decompose requires a valid-dimensional vector field") # raise ValueError("nHHD.decompose requires a valid-dimensional vector field")
if verbose > 0: # if verbose > 0:
print( '\n ==> nhhd.decompose', vfield.shape) # print( '\n ==> nhhd.decompose', vfield.shape)
if self.dim == 2: # if self.dim == 2:
# compute div curl using GPU # # compute div curl using GPU
(self.div, self.curlw) = self.mesh.divcurl(vfield) # (self.div, self.curlw) = self.mesh.divcurl(vfield)
if verbose > 0: # if verbose > 0:
print( 'div =', self.div.shape, self.div.min(), self.div.max()) # print( 'div =', self.div.shape, self.div.min(), self.div.max())
print( 'curlw =', self.curlw.shape, self.curlw.min(), self.curlw.max()) # print( 'curlw =', self.curlw.shape, self.curlw.min(), self.curlw.max())
# compute potentials using batch GPU solver # # compute potentials using batch GPU solver
batch_inputs = [self.div, self.curlw] # batch_inputs = [self.div, self.curlw]
batch_results = self.psolver.solve_gpu_batch(batch_inputs, verbose > 0) # batch_results = self.psolver.solve_gpu_batch(batch_inputs, verbose > 0)
self.nD = batch_results[0] # self.nD = batch_results[0]
self.nRu = batch_results[1] # self.nRu = batch_results[1]
if verbose > 0: # if verbose > 0:
print( 'nD =', self.nD.shape, self.nD.min(), self.nD.max()) # print( 'nD =', self.nD.shape, self.nD.min(), self.nD.max())
print( 'nR =', self.nRu.shape, self.nRu.min(), self.nRu.max()) # print( 'nR =', self.nRu.shape, self.nRu.min(), self.nRu.max())
# compute fields as gradients of potentials # # compute fields as gradients of potentials
self.d = self.mesh.gradient(self.nD) # self.d = self.mesh.gradient(self.nD)
if verbose > 0: # if verbose > 0:
print( 'd =', self.d.shape, self.d.min(), self.d.max()) # print( 'd =', self.d.shape, self.d.min(), self.d.max())
self.r = self.mesh.rotated_gradient(self.nRu) # self.r = self.mesh.rotated_gradient(self.nRu)
if verbose > 0: # if verbose > 0:
print( 'r =', self.r.shape, self.r.min(), self.r.max()) # print( 'r =', self.r.shape, self.r.min(), self.r.max())
else: # else:
# compute div curl # # compute div curl
(self.div, self.curlu, self.curlv, self.curlw) = self.mesh.divcurl(vfield) # (self.div, self.curlu, self.curlv, self.curlw) = self.mesh.divcurl(vfield)
if verbose > 0: # if verbose > 0:
print( 'div =', self.div.shape, self.div.min(), self.div.max()) # print( 'div =', self.div.shape, self.div.min(), self.div.max())
print( 'curlu =', self.curlu.shape, self.curlu.min(), self.curlu.max()) # print( 'curlu =', self.curlu.shape, self.curlu.min(), self.curlu.max())
print( 'curlv =', self.curlv.shape, self.curlv.min(), self.curlv.max()) # print( 'curlv =', self.curlv.shape, self.curlv.min(), self.curlv.max())
print( 'curlw =', self.curlw.shape, self.curlw.min(), self.curlw.max()) # print( 'curlw =', self.curlw.shape, self.curlw.min(), self.curlw.max())
# compute potentials using batch GPU solver # # compute potentials using batch GPU solver
batch_inputs = [self.div, self.curlu, self.curlv, self.curlw] # batch_inputs = [self.div, self.curlu, self.curlv, self.curlw]
batch_results = self.psolver.solve_gpu_batch(batch_inputs, verbose > 0) # batch_results = self.psolver.solve_gpu_batch(batch_inputs, verbose > 0)
self.nD = batch_results[0] # self.nD = batch_results[0]
self.nRu = batch_results[1] * -1.0 # self.nRu = batch_results[1] * -1.0
self.nRv = batch_results[2] * -1.0 # self.nRv = batch_results[2] * -1.0
self.nRw = batch_results[3] * -1.0 # self.nRw = batch_results[3] * -1.0
if verbose > 0: # if verbose > 0:
print( 'nD =', self.nD.shape, self.nD.min(), self.nD.max()) # print( 'nD =', self.nD.shape, self.nD.min(), self.nD.max())
print( 'nRu =', self.nRu.shape, self.nRu.min(), self.nRu.max()) # print( 'nRu =', self.nRu.shape, self.nRu.min(), self.nRu.max())
print( 'nRv =', self.nRv.shape, self.nRv.min(), self.nRv.max()) # print( 'nRv =', self.nRv.shape, self.nRv.min(), self.nRv.max())
print( 'nRw =', self.nRw.shape, self.nRw.min(), self.nRw.max()) # print( 'nRw =', self.nRw.shape, self.nRw.min(), self.nRw.max())
# compute divergent field as gradient of D potential # # compute divergent field as gradient of D potential
self.d = self.mesh.gradient(self.nD) # self.d = self.mesh.gradient(self.nD)
if verbose > 0: # if verbose > 0:
print( 'd =', self.d.shape, self.d.min(), self.d.max()) # print( 'd =', self.d.shape, self.d.min(), self.d.max())
# temporary use of r # # temporary use of r
self.r = numpy.stack((self.nRu, self.nRv, self.nRw), axis = -1) # self.r = numpy.stack((self.nRu, self.nRv, self.nRw), axis = -1)
# compute rotational field as curl of R potential # # compute rotational field as curl of R potential
(cu, cv, cw) = self.mesh.curl3D(self.r) # (cu, cv, cw) = self.mesh.curl3D(self.r)
self.r = numpy.stack((cu, cv, cw), axis = -1) # self.r = numpy.stack((cu, cv, cw), axis = -1)
if verbose > 0: # if verbose > 0:
print( 'r =', self.r.shape, self.r.min(), self.r.max()) # print( 'r =', self.r.shape, self.r.min(), self.r.max())
self.h = numpy.add(self.d, self.r) # self.h = numpy.add(self.d, self.r)
numpy.subtract(vfield, self.h, self.h) # numpy.subtract(vfield, self.h, self.h)
if verbose > 0: # if verbose > 0:
print( 'h =', self.h.shape, self.h.min(), self.h.max()) # print( 'h =', self.h.shape, self.h.min(), self.h.max())
def get_force(self, vfield, verbose=False): def get_force(self, vfield, verbose=False):
@@ -434,89 +434,89 @@ class nHHD(object):
self._force_gpu_cache['mgrid_1'], shape self._force_gpu_cache['mgrid_1'], shape
).copy() ).copy()
def get_force_gpu(self, vfield, verbose=False): # def get_force_gpu(self, vfield, verbose=False):
""" # """
GPU-accelerated force computation using CuPy. # GPU-accelerated force computation using CuPy.
Falls back to CPU if CuPy is not available. # Falls back to CPU if CuPy is not available.
Args: # Args:
vfield: Vector field (numpy array) # vfield: Vector field (numpy array)
verbose: Print debug information (default: False) # verbose: Print debug information (default: False)
Returns: # Returns:
tuple: (flag, output) where: # tuple: (flag, output) where:
- flag: Boolean indicating if force thresholds are met # - flag: Boolean indicating if force thresholds are met
- output: Array [Sn, Stx, Sty, Stor] containing force components # - output: Array [Sn, Stx, Sty, Stor] containing force components
""" # """
if not CUPY_AVAILABLE: # if not CUPY_AVAILABLE:
if verbose: # if verbose:
print("CuPy not available, falling back to CPU get_force") # print("CuPy not available, falling back to CPU get_force")
return self.get_force(vfield, verbose) # return self.get_force(vfield, verbose)
# Decompose the vector field (uses GPU internally) # # Decompose the vector field (uses GPU internally)
self.decompose(vfield, verbose) # self.decompose(vfield, verbose)
# Initialize cache if needed or shape changed # # Initialize cache if needed or shape changed
shape = (vfield.shape[0], vfield.shape[1]) # shape = (vfield.shape[0], vfield.shape[1])
if not hasattr(self, '_force_gpu_cache') or self._force_gpu_cache is None or self._force_gpu_cache['shape'] != shape: # 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) # self._init_force_gpu_cache(shape)
# Transfer data to GPU (consider keeping on GPU in decompose for further optimization) # # Transfer data to GPU (consider keeping on GPU in decompose for further optimization)
d_gpu = cp.asarray(self.d, dtype=cp.float32) # d_gpu = cp.asarray(self.d, dtype=cp.float32)
h_gpu = cp.asarray(self.h, dtype=cp.float32) # h_gpu = cp.asarray(self.h, dtype=cp.float32)
r_gpu = cp.asarray(self.r, dtype=cp.float32) # r_gpu = cp.asarray(self.r, dtype=cp.float32)
nRu_gpu = cp.asarray(self.nRu, dtype=cp.float32) # nRu_gpu = cp.asarray(self.nRu, dtype=cp.float32)
# Compute all sums in parallel - fused for better performance # # Compute all sums in parallel - fused for better performance
# Sn: sum of absolute values of divergent field (both components) # # Sn: sum of absolute values of divergent field (both components)
Sn_gpu = cp.sum(cp.abs(d_gpu)) # Sn_gpu = cp.sum(cp.abs(d_gpu))
# Stx and Sty: sum of harmonic field components # # Stx and Sty: sum of harmonic field components
Stx_gpu = cp.sum(h_gpu[:, :, 0]) # Stx_gpu = cp.sum(h_gpu[:, :, 0])
Sty_gpu = cp.sum(h_gpu[:, :, 1]) # Sty_gpu = cp.sum(h_gpu[:, :, 1])
# Find max/min locations of nRu using single pass # # Find max/min locations of nRu using single pass
nR_max_idx = int(cp.argmax(nRu_gpu)) # nR_max_idx = int(cp.argmax(nRu_gpu))
nR_min_idx = int(cp.argmin(nRu_gpu)) # nR_min_idx = int(cp.argmin(nRu_gpu))
# Convert flat indices to 2D coordinates (on CPU - faster for single values) # # Convert flat indices to 2D coordinates (on CPU - faster for single values)
nR_max_loc = numpy.unravel_index(nR_max_idx, nRu_gpu.shape) # nR_max_loc = numpy.unravel_index(nR_max_idx, nRu_gpu.shape)
nR_min_loc = numpy.unravel_index(nR_min_idx, nRu_gpu.shape) # nR_min_loc = numpy.unravel_index(nR_min_idx, nRu_gpu.shape)
# Use cached grid arrays # # Use cached grid arrays
mgrid_0 = self._force_gpu_cache['mgrid_0_bc'] # mgrid_0 = self._force_gpu_cache['mgrid_0_bc']
mgrid_1 = self._force_gpu_cache['mgrid_1_bc'] # mgrid_1 = self._force_gpu_cache['mgrid_1_bc']
# Compute torque using fused operations # # Compute torque using fused operations
# r_p and r_n displacements + cross products in one pass # # r_p and r_n displacements + cross products in one pass
max_0, max_1 = float(nR_max_loc[0]), float(nR_max_loc[1]) # max_0, max_1 = float(nR_max_loc[0]), float(nR_max_loc[1])
min_0, min_1 = float(nR_min_loc[0]), float(nR_min_loc[1]) # min_0, min_1 = float(nR_min_loc[0]), float(nR_min_loc[1])
r_x = r_gpu[:, :, 0] # r_x = r_gpu[:, :, 0]
r_y = r_gpu[:, :, 1] # r_y = r_gpu[:, :, 1]
# Fused cross product computation # # Fused cross product computation
# cross_p = (mgrid_0 - max_0) * r_y - (mgrid_1 - max_1) * r_x # # cross_p = (mgrid_0 - max_0) * r_y - (mgrid_1 - max_1) * r_x
# cross_n = (mgrid_0 - min_0) * r_y - (mgrid_1 - min_1) * r_x # # cross_n = (mgrid_0 - min_0) * r_y - (mgrid_1 - min_1) * r_x
Stor_gpu = cp.sum( # Stor_gpu = cp.sum(
(mgrid_0 - max_0) * r_y - (mgrid_1 - max_1) * r_x + # (mgrid_0 - max_0) * r_y - (mgrid_1 - max_1) * r_x +
(mgrid_0 - min_0) * r_y - (mgrid_1 - min_1) * r_x # (mgrid_0 - min_0) * r_y - (mgrid_1 - min_1) * r_x
) # )
# Single batch transfer to CPU (more efficient than 4 separate transfers) # # Single batch transfer to CPU (more efficient than 4 separate transfers)
results = cp.asnumpy(cp.array([Sn_gpu, Stx_gpu, Sty_gpu, Stor_gpu])) # results = cp.asnumpy(cp.array([Sn_gpu, Stx_gpu, Sty_gpu, Stor_gpu]))
self.Sn = float(results[0]) # self.Sn = float(results[0])
self.Stx = float(results[1]) # self.Stx = float(results[1])
self.Sty = float(results[2]) # self.Sty = float(results[2])
self.Stor = float(results[3]) # self.Stor = float(results[3])
output = results # output = results
# Check thresholds # # Check thresholds
flag = (self.Sn >= 300 and self.Stx >= 300 and self.Sty >= 300 and numpy.abs(self.Stor) >= 1000) # flag = (self.Sn >= 300 and self.Stx >= 300 and self.Sty >= 300 and numpy.abs(self.Stor) >= 1000)
return flag, output # return flag, output
def get_force_gpu_v2(self, vfield, verbose=False): def get_force_gpu_v2(self, vfield, verbose=False):
""" """
@@ -592,190 +592,190 @@ class nHHD(object):
if CUPY_AVAILABLE: if CUPY_AVAILABLE:
cp.get_default_memory_pool().free_all_blocks() cp.get_default_memory_pool().free_all_blocks()
def get_force_gpu_stream(self, vfield, verbose=False): # def get_force_gpu_stream(self, vfield, verbose=False):
""" # """
Fully optimized GPU force computation using: # Fully optimized GPU force computation using:
- decompose_gpu to keep data on GPU throughout # - decompose_gpu to keep data on GPU throughout
- CUDA streams for parallel reductions # - CUDA streams for parallel reductions
- Fused kernel operations # - Fused kernel operations
Args: # Args:
vfield: Vector field (numpy array) # vfield: Vector field (numpy array)
verbose: Print debug information (default: False) # verbose: Print debug information (default: False)
Returns: # Returns:
tuple: (flag, output) where: # tuple: (flag, output) where:
- flag: Boolean indicating if force thresholds are met # - flag: Boolean indicating if force thresholds are met
- output: Array [Sn, Stx, Sty, Stor] containing force components # - output: Array [Sn, Stx, Sty, Stor] containing force components
""" # """
if not CUPY_AVAILABLE: # if not CUPY_AVAILABLE:
if verbose: # if verbose:
print("CuPy not available, falling back to CPU get_force") # print("CuPy not available, falling back to CPU get_force")
return self.get_force(vfield, verbose) # return self.get_force(vfield, verbose)
# Use GPU decomposition to keep data on GPU # # Use GPU decomposition to keep data on GPU
gpu_result = self.decompose_gpu(vfield, verbose) # gpu_result = self.decompose_gpu(vfield, verbose)
if gpu_result is None: # if gpu_result is None:
# Fallback for 3D or if decompose_gpu failed # # Fallback for 3D or if decompose_gpu failed
return self.get_force_gpu(vfield, verbose) # return self.get_force_gpu(vfield, verbose)
d_gpu, h_gpu, r_gpu, nRu_gpu = gpu_result # d_gpu, h_gpu, r_gpu, nRu_gpu = gpu_result
# Initialize cache if needed or shape changed # # Initialize cache if needed or shape changed
shape = (vfield.shape[0], vfield.shape[1]) # shape = (vfield.shape[0], vfield.shape[1])
if not hasattr(self, '_force_gpu_cache') or self._force_gpu_cache is None or self._force_gpu_cache['shape'] != shape: # 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) # self._init_force_gpu_cache(shape)
# Use CUDA streams for parallel computation if available # # Use CUDA streams for parallel computation if available
if STREAMS_AVAILABLE: # if STREAMS_AVAILABLE:
stream1 = Stream(non_blocking=True) # stream1 = Stream(non_blocking=True)
stream2 = Stream(non_blocking=True) # stream2 = Stream(non_blocking=True)
stream3 = Stream(non_blocking=True) # stream3 = Stream(non_blocking=True)
# Stream 1: Compute Sn # # Stream 1: Compute Sn
with stream1: # with stream1:
Sn_gpu = cp.sum(cp.abs(d_gpu)) # Sn_gpu = cp.sum(cp.abs(d_gpu))
# Stream 2: Compute Stx, Sty # # Stream 2: Compute Stx, Sty
with stream2: # with stream2:
Stx_gpu = cp.sum(h_gpu[:, :, 0]) # Stx_gpu = cp.sum(h_gpu[:, :, 0])
Sty_gpu = cp.sum(h_gpu[:, :, 1]) # Sty_gpu = cp.sum(h_gpu[:, :, 1])
# Stream 3: Find extrema and compute torque # # Stream 3: Find extrema and compute torque
with stream3: # with stream3:
nR_max_idx = int(cp.argmax(nRu_gpu)) # nR_max_idx = int(cp.argmax(nRu_gpu))
nR_min_idx = int(cp.argmin(nRu_gpu)) # nR_min_idx = int(cp.argmin(nRu_gpu))
max_0, max_1 = nR_max_idx // shape[1], nR_max_idx % shape[1] # 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] # min_0, min_1 = nR_min_idx // shape[1], nR_min_idx % shape[1]
mgrid_0 = self._force_gpu_cache['mgrid_0_bc'] # mgrid_0 = self._force_gpu_cache['mgrid_0_bc']
mgrid_1 = self._force_gpu_cache['mgrid_1_bc'] # mgrid_1 = self._force_gpu_cache['mgrid_1_bc']
r_x = r_gpu[:, :, 0] # r_x = r_gpu[:, :, 0]
r_y = r_gpu[:, :, 1] # r_y = r_gpu[:, :, 1]
Stor_gpu = cp.sum( # Stor_gpu = cp.sum(
(mgrid_0 - max_0) * r_y - (mgrid_1 - max_1) * r_x + # (mgrid_0 - max_0) * r_y - (mgrid_1 - max_1) * r_x +
(mgrid_0 - min_0) * r_y - (mgrid_1 - min_1) * r_x # (mgrid_0 - min_0) * r_y - (mgrid_1 - min_1) * r_x
) # )
# Synchronize all streams # # Synchronize all streams
stream1.synchronize() # stream1.synchronize()
stream2.synchronize() # stream2.synchronize()
stream3.synchronize() # stream3.synchronize()
else: # else:
# Without streams, compute sequentially but still on GPU # # Without streams, compute sequentially but still on GPU
Sn_gpu = cp.sum(cp.abs(d_gpu)) # Sn_gpu = cp.sum(cp.abs(d_gpu))
Stx_gpu = cp.sum(h_gpu[:, :, 0]) # Stx_gpu = cp.sum(h_gpu[:, :, 0])
Sty_gpu = cp.sum(h_gpu[:, :, 1]) # Sty_gpu = cp.sum(h_gpu[:, :, 1])
nR_max_idx = int(cp.argmax(nRu_gpu)) # nR_max_idx = int(cp.argmax(nRu_gpu))
nR_min_idx = int(cp.argmin(nRu_gpu)) # nR_min_idx = int(cp.argmin(nRu_gpu))
max_0, max_1 = nR_max_idx // shape[1], nR_max_idx % shape[1] # 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] # min_0, min_1 = nR_min_idx // shape[1], nR_min_idx % shape[1]
mgrid_0 = self._force_gpu_cache['mgrid_0_bc'] # mgrid_0 = self._force_gpu_cache['mgrid_0_bc']
mgrid_1 = self._force_gpu_cache['mgrid_1_bc'] # mgrid_1 = self._force_gpu_cache['mgrid_1_bc']
r_x = r_gpu[:, :, 0] # r_x = r_gpu[:, :, 0]
r_y = r_gpu[:, :, 1] # r_y = r_gpu[:, :, 1]
Stor_gpu = cp.sum( # Stor_gpu = cp.sum(
(mgrid_0 - max_0) * r_y - (mgrid_1 - max_1) * r_x + # (mgrid_0 - max_0) * r_y - (mgrid_1 - max_1) * r_x +
(mgrid_0 - min_0) * r_y - (mgrid_1 - min_1) * r_x # (mgrid_0 - min_0) * r_y - (mgrid_1 - min_1) * r_x
) # )
# Single batch transfer to CPU # # Single batch transfer to CPU
results = cp.asnumpy(cp.array([Sn_gpu, Stx_gpu, Sty_gpu, Stor_gpu])) # results = cp.asnumpy(cp.array([Sn_gpu, Stx_gpu, Sty_gpu, Stor_gpu]))
self.Sn = float(results[0]) # self.Sn = float(results[0])
self.Stx = float(results[1]) # self.Stx = float(results[1])
self.Sty = float(results[2]) # self.Sty = float(results[2])
self.Stor = float(results[3]) # self.Stor = float(results[3])
output = results # output = results
# Check thresholds # # Check thresholds
flag = (self.Sn >= 300 and self.Stx >= 300 and self.Sty >= 300 and numpy.abs(self.Stor) >= 1000) # flag = (self.Sn >= 300 and self.Stx >= 300 and self.Sty >= 300 and numpy.abs(self.Stor) >= 1000)
return flag, output # return flag, output
def get_force_gpu_kernel(self, vfield, verbose=False): # def get_force_gpu_kernel(self, vfield, verbose=False):
""" # """
Maximum performance GPU force computation using: # Maximum performance GPU force computation using:
- decompose_gpu to keep data on GPU throughout # - decompose_gpu to keep data on GPU throughout
- Custom CUDA kernel for fused torque computation # - Custom CUDA kernel for fused torque computation
- Minimized memory allocations # - Minimized memory allocations
Args: # Args:
vfield: Vector field (numpy array) # vfield: Vector field (numpy array)
verbose: Print debug information (default: False) # verbose: Print debug information (default: False)
Returns: # Returns:
tuple: (flag, output) where: # tuple: (flag, output) where:
- flag: Boolean indicating if force thresholds are met # - flag: Boolean indicating if force thresholds are met
- output: Array [Sn, Stx, Sty, Stor] containing force components # - output: Array [Sn, Stx, Sty, Stor] containing force components
""" # """
if not CUPY_AVAILABLE: # if not CUPY_AVAILABLE:
if verbose: # if verbose:
print("CuPy not available, falling back to CPU get_force") # print("CuPy not available, falling back to CPU get_force")
return self.get_force(vfield, verbose) # return self.get_force(vfield, verbose)
# Use GPU decomposition to keep data on GPU # # Use GPU decomposition to keep data on GPU
gpu_result = self.decompose_gpu(vfield, verbose) # gpu_result = self.decompose_gpu(vfield, verbose)
if gpu_result is None: # if gpu_result is None:
return self.get_force_gpu(vfield, verbose) # return self.get_force_gpu(vfield, verbose)
d_gpu, h_gpu, r_gpu, nRu_gpu = gpu_result # d_gpu, h_gpu, r_gpu, nRu_gpu = gpu_result
# Initialize cache if needed # # Initialize cache if needed
shape = (vfield.shape[0], vfield.shape[1]) # shape = (vfield.shape[0], vfield.shape[1])
if not hasattr(self, '_force_gpu_cache') or self._force_gpu_cache is None or self._force_gpu_cache['shape'] != shape: # 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) # self._init_force_gpu_cache(shape)
# Compute Sn using L1 norm (faster than sum(abs())) # # Compute Sn using L1 norm (faster than sum(abs()))
Sn = float(cp.linalg.norm(d_gpu.ravel(), ord=1)) # Sn = float(cp.linalg.norm(d_gpu.ravel(), ord=1))
# Compute Stx, Sty # # Compute Stx, Sty
Stx = float(cp.sum(h_gpu[:, :, 0])) # Stx = float(cp.sum(h_gpu[:, :, 0]))
Sty = float(cp.sum(h_gpu[:, :, 1])) # Sty = float(cp.sum(h_gpu[:, :, 1]))
# Find extrema # # Find extrema
nR_max_idx = int(cp.argmax(nRu_gpu)) # nR_max_idx = int(cp.argmax(nRu_gpu))
nR_min_idx = int(cp.argmin(nRu_gpu)) # nR_min_idx = int(cp.argmin(nRu_gpu))
max_0 = cp.float32(nR_max_idx // shape[1]) # max_0 = cp.float32(nR_max_idx // shape[1])
max_1 = cp.float32(nR_max_idx % shape[1]) # max_1 = cp.float32(nR_max_idx % shape[1])
min_0 = cp.float32(nR_min_idx // shape[1]) # min_0 = cp.float32(nR_min_idx // shape[1])
min_1 = cp.float32(nR_min_idx % shape[1]) # min_1 = cp.float32(nR_min_idx % shape[1])
mgrid_0 = self._force_gpu_cache['mgrid_0_bc'] # mgrid_0 = self._force_gpu_cache['mgrid_0_bc']
mgrid_1 = self._force_gpu_cache['mgrid_1_bc'] # mgrid_1 = self._force_gpu_cache['mgrid_1_bc']
r_x = r_gpu[:, :, 0] # r_x = r_gpu[:, :, 0]
r_y = r_gpu[:, :, 1] # r_y = r_gpu[:, :, 1]
# Use custom kernel if available for fused torque computation # # Use custom kernel if available for fused torque computation
if KERNEL_AVAILABLE and _torque_kernel is not None: # if KERNEL_AVAILABLE and _torque_kernel is not None:
torque_vals = _torque_kernel(r_x, r_y, mgrid_0, mgrid_1, max_0, max_1, min_0, min_1) # torque_vals = _torque_kernel(r_x, r_y, mgrid_0, mgrid_1, max_0, max_1, min_0, min_1)
Stor = float(cp.sum(torque_vals)) # Stor = float(cp.sum(torque_vals))
else: # else:
# Fallback to standard computation # # Fallback to standard computation
Stor = float(cp.sum( # Stor = float(cp.sum(
(mgrid_0 - max_0) * r_y - (mgrid_1 - max_1) * r_x + # (mgrid_0 - max_0) * r_y - (mgrid_1 - max_1) * r_x +
(mgrid_0 - min_0) * r_y - (mgrid_1 - min_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 # self.Sn, self.Stx, self.Sty, self.Stor = Sn, Stx, Sty, Stor
output = numpy.array([Sn, Stx, Sty, Stor], dtype=numpy.float32) # output = numpy.array([Sn, Stx, Sty, Stor], dtype=numpy.float32)
flag = (Sn >= 300 and Stx >= 300 and Sty >= 300 and abs(Stor) >= 1000) # flag = (Sn >= 300 and Stx >= 300 and Sty >= 300 and abs(Stor) >= 1000)
return flag, output # return flag, output
def get_shear(self): def get_shear(self):
src/orisys/sensor.py
+1 -1
View File
@@ -17,7 +17,7 @@ from queue import Queue, Empty
from .Stitcher import Stitcher from .Stitcher import Stitcher
from .tracker import Tracker from .tracker import Tracker
from .nHHD import nHHD from .HHD import nHHD
from . import configs from . import configs
from .get_depth import Contact3D from .get_depth import Contact3D