combined HHD
This commit is contained in:
bluepulsar
@@ -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()
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# --------------------------------------------------------------------------
|
||||
Reference in New Issue
Block a user