nuitka build success to __init__.pyd
This commit is contained in:
bluepulsar
+1
-1
@@ -860,7 +860,7 @@ class StructuredGrid(object):
|
||||
|
||||
import numpy
|
||||
from scipy import signal, spatial
|
||||
from .timer import Timer
|
||||
# from .timer import Timer
|
||||
|
||||
# Try to import CuPy for GPU acceleration
|
||||
try:
|
||||
|
||||
@@ -1,315 +0,0 @@
|
||||
"""
|
||||
MultiSensor - Multi-process sensor management for parallel processing.
|
||||
|
||||
This module provides the MultiSensor class for running multiple Sensor
|
||||
instances in parallel processes, enabling true parallel processing with
|
||||
independent CUDA contexts.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
|
||||
def _sensor_worker(sensor_id, config, data_queue, stop_event):
|
||||
"""
|
||||
Worker function that runs in a separate PROCESS for each sensor.
|
||||
Each process has its own Python interpreter and CUDA context.
|
||||
"""
|
||||
os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0"
|
||||
|
||||
from .sensor import Sensor
|
||||
|
||||
print(f"[Sensor {sensor_id}] Process started (PID: {os.getpid()})")
|
||||
|
||||
sensor = Sensor(
|
||||
vid_src=config.get('vid_src'),
|
||||
isstitch=config.get('isstitch', True),
|
||||
cuda=config.get('cuda', True),
|
||||
verbose=config.get('verbose', False),
|
||||
config_name=config.get('config_name', 'prototype'),
|
||||
cal_path=config.get('cal_path', None)
|
||||
)
|
||||
compute_contact = config.get('compute_contact', False)
|
||||
|
||||
frame_count = 0
|
||||
|
||||
def queue_put(data):
|
||||
try:
|
||||
while not data_queue.empty():
|
||||
try: data_queue.get_nowait()
|
||||
except: break
|
||||
data_queue.put_nowait(data)
|
||||
except: pass
|
||||
|
||||
while not stop_event.is_set():
|
||||
try:
|
||||
# Sync mode
|
||||
sensor.get_img()
|
||||
sensor.compute_deformation()
|
||||
if compute_contact:
|
||||
sensor.compute_contact()
|
||||
|
||||
# Use read_info to get all sensor data
|
||||
|
||||
img, flow, normal, shear, fps, fnormal, fshearx, fsheary = sensor.read_info(
|
||||
sensor.info.IMG,
|
||||
sensor.info.VRAW,
|
||||
sensor.info.VNORMAL,
|
||||
sensor.info.VSHEAR,
|
||||
sensor.info.FPS,
|
||||
sensor.info.FNORMAL,
|
||||
sensor.info.FSHEARX,
|
||||
sensor.info.FSHEARY
|
||||
)
|
||||
if compute_contact:
|
||||
depth, centroid, contour = sensor.read_info(
|
||||
sensor.info.DEPTH,
|
||||
sensor.info.CENTROID,
|
||||
sensor.info.CONTOUR
|
||||
)
|
||||
else:
|
||||
depth = None
|
||||
centroid = None
|
||||
contour = None
|
||||
|
||||
# Build data dictionary
|
||||
data = {
|
||||
'sensor_id': sensor_id,
|
||||
'img': img if img is not None and img.size > 0 else None,
|
||||
'flow': flow if flow is not None and flow.size > 0 else None,
|
||||
'normal': normal if normal is not None and normal.size > 0 else None,
|
||||
'shear': shear if shear is not None and shear.size > 0 else None,
|
||||
'fps': fps,
|
||||
'fnormal': fnormal,
|
||||
'fshearx': fshearx,
|
||||
'fsheary': fsheary,
|
||||
'frame_count': frame_count,
|
||||
}
|
||||
|
||||
# Handle contact-related fields
|
||||
if compute_contact:
|
||||
# depth and depth_map
|
||||
data['depth'] = depth if depth is not None and depth.size > 0 else None
|
||||
data['depth_map'] = data['depth'] # Alias for compatibility
|
||||
|
||||
# centroid - can be None or numpy array
|
||||
if centroid is not None:
|
||||
# Check if it's a valid array with elements
|
||||
if hasattr(centroid, 'size') and centroid.size > 0:
|
||||
data['centroid'] = centroid
|
||||
data['centroid_found'] = True
|
||||
else:
|
||||
data['centroid'] = None
|
||||
data['centroid_found'] = False
|
||||
else:
|
||||
data['centroid'] = None
|
||||
data['centroid_found'] = False
|
||||
|
||||
# contour
|
||||
data['contour'] = contour if contour is not None and contour.size > 0 else None
|
||||
else:
|
||||
data['depth'] = None
|
||||
data['depth_map'] = None
|
||||
data['centroid'] = None
|
||||
data['centroid_found'] = False
|
||||
data['contour'] = None
|
||||
|
||||
queue_put(data)
|
||||
frame_count += 1
|
||||
|
||||
except Exception as e:
|
||||
print(f"[Sensor {sensor_id}] Error: {e}")
|
||||
break
|
||||
|
||||
sensor.disconnect()
|
||||
print(f"[Sensor {sensor_id}] Process stopped")
|
||||
|
||||
|
||||
class MultiSensor:
|
||||
"""
|
||||
Manages multiple Sensor instances running in parallel processes.
|
||||
|
||||
Each sensor runs in its own process with independent CUDA context,
|
||||
enabling true parallel processing without GIL limitations.
|
||||
|
||||
Example:
|
||||
# Simple usage with context manager
|
||||
configs = [
|
||||
{'vid_src': './s1.mp4', 'isstitch': True},
|
||||
{'vid_src': './s2.mp4', 'isstitch': True},
|
||||
]
|
||||
|
||||
with MultiSensor(configs) as multi:
|
||||
for all_data in multi:
|
||||
for data in all_data:
|
||||
if data:
|
||||
print(f"Sensor {data['sensor_id']}: FPS={data['fps']}")
|
||||
cv2.imshow(f"Sensor {data['sensor_id']}", data['img'])
|
||||
if cv2.waitKey(1) & 0xFF == ord('q'):
|
||||
break
|
||||
|
||||
# Or manual control
|
||||
multi = MultiSensor(configs)
|
||||
multi.start()
|
||||
|
||||
data = multi.get(0) # Get from sensor 0
|
||||
all_data = multi.get_all() # Get from all sensors
|
||||
|
||||
multi.stop()
|
||||
"""
|
||||
|
||||
def __init__(self, sensor_configs):
|
||||
"""
|
||||
Initialize MultiSensor manager.
|
||||
|
||||
Args:
|
||||
sensor_configs: List of dicts, each containing Sensor constructor args:
|
||||
- vid_src: Video source (required) - int for camera ID, str for file path
|
||||
- isstitch: Whether to stitch (default: True)
|
||||
- cuda: Use CUDA acceleration (default: True)
|
||||
- verbose: Print logs (default: False)
|
||||
- config_name: Config name (default: 'prototype')
|
||||
- cal_path: Calibration file path (default: None)
|
||||
|
||||
Example:
|
||||
configs = [
|
||||
{'vid_src': 0, 'isstitch': True, 'cuda': True},
|
||||
{'vid_src': 1, 'isstitch': True, 'cuda': True},
|
||||
]
|
||||
multi = MultiSensor(configs)
|
||||
"""
|
||||
import multiprocessing as mp
|
||||
|
||||
self.sensor_configs = sensor_configs
|
||||
self.num_sensors = len(sensor_configs)
|
||||
self._mp = mp
|
||||
|
||||
self.processes = []
|
||||
self.data_queues = []
|
||||
self.stop_events = []
|
||||
self._started = False
|
||||
|
||||
# Create queues and events for each sensor
|
||||
for i in range(self.num_sensors):
|
||||
self.data_queues.append(mp.Queue(maxsize=2))
|
||||
self.stop_events.append(mp.Event())
|
||||
|
||||
def start(self):
|
||||
"""Start all sensor worker processes."""
|
||||
if self._started:
|
||||
print("MultiSensor already started")
|
||||
return
|
||||
|
||||
for i, config in enumerate(self.sensor_configs):
|
||||
process = self._mp.Process(
|
||||
target=_sensor_worker,
|
||||
args=(i, config, self.data_queues[i], self.stop_events[i]),
|
||||
daemon=True
|
||||
)
|
||||
process.start()
|
||||
self.processes.append(process)
|
||||
print(f"Started sensor {i} (PID: {process.pid})")
|
||||
|
||||
self._started = True
|
||||
print(f"All {len(self.processes)} sensors started")
|
||||
|
||||
def get(self, sensor_idx, timeout=0.01):
|
||||
"""
|
||||
Get the latest data from a specific sensor.
|
||||
|
||||
Args:
|
||||
sensor_idx: Index of the sensor (0-based)
|
||||
timeout: How long to wait for data (seconds)
|
||||
|
||||
Returns:
|
||||
dict: Sensor data or None if no data available
|
||||
- sensor_id: int
|
||||
- img: np.ndarray (H, W, 3)
|
||||
- flow: np.ndarray (H, W, 2)
|
||||
- normal: np.ndarray
|
||||
- shear: np.ndarray
|
||||
- depth: np.ndarray
|
||||
- fps: float
|
||||
- fnormal, fshearx, fsheary, ftor: float
|
||||
- frame_count: int
|
||||
"""
|
||||
try:
|
||||
return self.data_queues[sensor_idx].get(timeout=timeout)
|
||||
except:
|
||||
return None
|
||||
|
||||
def get_all(self, timeout=0.01):
|
||||
"""
|
||||
Get the latest data from all sensors.
|
||||
|
||||
Returns:
|
||||
list: List of sensor data dicts (None for sensors with no new data)
|
||||
"""
|
||||
return [self.get(i, timeout) for i in range(self.num_sensors)]
|
||||
|
||||
def __iter__(self):
|
||||
"""
|
||||
Iterate over sensor data. Yields all sensor data on each iteration.
|
||||
|
||||
Example:
|
||||
for all_data in multi:
|
||||
for data in all_data:
|
||||
if data:
|
||||
process(data)
|
||||
"""
|
||||
return self
|
||||
|
||||
def __next__(self):
|
||||
"""Get next batch of sensor data."""
|
||||
if not self._started:
|
||||
raise StopIteration("MultiSensor not started")
|
||||
|
||||
# Check if any process is still alive
|
||||
if not any(p.is_alive() for p in self.processes):
|
||||
raise StopIteration("All sensor processes have stopped")
|
||||
|
||||
return self.get_all()
|
||||
|
||||
def __len__(self):
|
||||
"""Return number of sensors."""
|
||||
return self.num_sensors
|
||||
|
||||
def __getitem__(self, idx):
|
||||
"""Get latest data from sensor by index."""
|
||||
return self.get(idx)
|
||||
|
||||
def stop(self):
|
||||
"""Stop all sensor processes gracefully."""
|
||||
if not self._started:
|
||||
return
|
||||
|
||||
print("Stopping all sensors...")
|
||||
|
||||
# Signal all processes to stop
|
||||
for event in self.stop_events:
|
||||
event.set()
|
||||
|
||||
# Wait for processes to finish
|
||||
for i, process in enumerate(self.processes):
|
||||
process.join(timeout=3.0)
|
||||
if process.is_alive():
|
||||
print(f"Force terminating sensor {i}")
|
||||
process.terminate()
|
||||
|
||||
self.processes.clear()
|
||||
self._started = False
|
||||
print("All sensors stopped.")
|
||||
|
||||
def __enter__(self):
|
||||
"""Context manager entry - starts sensors."""
|
||||
self.start()
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
"""Context manager exit - stops sensors."""
|
||||
self.stop()
|
||||
return False
|
||||
|
||||
def __del__(self):
|
||||
"""Cleanup on deletion."""
|
||||
self.stop()
|
||||
|
||||
@@ -1,408 +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 sys
|
||||
from .timer import Timer
|
||||
|
||||
# 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()
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# --------------------------------------------------------------------------
|
||||
@@ -1,426 +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 sys
|
||||
import numpy
|
||||
from scipy import spatial
|
||||
from .timer import Timer
|
||||
|
||||
class UnstructuredGrid(object):
|
||||
|
||||
'''Class to support nHHD on unstructured grids (triangular and tetrahedral)'''
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# --------------------------------------------------------------------------
|
||||
def need_volumes(self, verbose=False):
|
||||
|
||||
''' Compute volumes/areas for vertices, simplices, and corners'''
|
||||
|
||||
if self.pvolumes.shape != (0,):
|
||||
return self.pvolumes
|
||||
|
||||
def tri_area(o,a,b):
|
||||
return numpy.linalg.norm(numpy.cross(b-o,a-o)) / 2.0
|
||||
|
||||
def tet_volume(o,a,b,c):
|
||||
return numpy.abs(numpy.dot(a-o,numpy.cross(b-o,c-o))) / 6.0
|
||||
|
||||
if verbose:
|
||||
print(' Computing point areas/volumes...', end ='')
|
||||
sys.stdout.flush()
|
||||
mtimer = Timer()
|
||||
|
||||
self.svolumes = numpy.zeros(self.nsimplices,)
|
||||
self.pvolumes = numpy.zeros(self.nvertices,)
|
||||
self.cvolumes = numpy.zeros((self.nsimplices,self.dim+1))
|
||||
|
||||
# triangulation
|
||||
if self.dim == 2:
|
||||
|
||||
for sidx in xrange(self.nsimplices):
|
||||
|
||||
simp = self.simplices[sidx]
|
||||
verts = self.vertices[simp]
|
||||
|
||||
self.svolumes[sidx] = tri_area(verts[0], verts[1], verts[2])
|
||||
|
||||
# edge weights
|
||||
e = [ verts[2]-verts[1], verts[0]-verts[2], verts[1]-verts[0] ]
|
||||
l2 = [ numpy.dot(e[0],e[0]), numpy.dot(e[1],e[1]), numpy.dot(e[2],e[2]) ]
|
||||
ew = [ l2[0]*(l2[1]+l2[2]-l2[0]), l2[1]*(l2[2]+l2[0]-l2[1]),l2[2]*(l2[0]+l2[1]-l2[2]) ]
|
||||
|
||||
# corner areas
|
||||
if (ew[0] <= 0):
|
||||
self.cvolumes[sidx,1] = -0.25 * l2[2] * self.svolumes[sidx] / numpy.dot(e[0], e[2])
|
||||
self.cvolumes[sidx,2] = -0.25 * l2[1] * self.svolumes[sidx] / numpy.dot(e[0], e[1])
|
||||
self.cvolumes[sidx,0] = self.svolumes[sidx] - self.cvolumes[sidx,1] - self.cvolumes[sidx,2]
|
||||
elif (ew[1] <= 0):
|
||||
self.cvolumes[sidx,2] = -0.25 * l2[0] * self.svolumes[sidx] / numpy.dot(e[1], e[0]);
|
||||
self.cvolumes[sidx,0] = -0.25 * l2[2] * self.svolumes[sidx] / numpy.dot(e[1], e[2]);
|
||||
self.cvolumes[sidx,1] = self.svolumes[sidx] - self.cvolumes[sidx,2] - self.cvolumes[sidx,0];
|
||||
elif (ew[2] <= 0):
|
||||
self.cvolumes[sidx,0] = -0.25 * l2[1] * self.svolumes[sidx] / numpy.dot(e[2], e[1]);
|
||||
self.cvolumes[sidx,1] = -0.25 * l2[0] * self.svolumes[sidx] / numpy.dot(e[2], e[0]);
|
||||
self.cvolumes[sidx,2] = self.svolumes[sidx] - self.cvolumes[sidx,0] - self.cvolumes[sidx,1];
|
||||
else:
|
||||
ewscale = 0.5 * self.svolumes[sidx] / (ew[0] + ew[1] + ew[2])
|
||||
for d in xrange(3):
|
||||
self.cvolumes[sidx,d] = ewscale * (ew[(d+1)%3] + ew[(d+2)%3])
|
||||
|
||||
self.pvolumes[simp[0]] += self.cvolumes[sidx,0]
|
||||
self.pvolumes[simp[1]] += self.cvolumes[sidx,1]
|
||||
self.pvolumes[simp[2]] += self.cvolumes[sidx,2]
|
||||
|
||||
# tetrahedralization
|
||||
elif self.sdim == 3:
|
||||
|
||||
raise ValueError('TODO: pvolumes for 3D')
|
||||
for sidx in xrange(self.nsimplices):
|
||||
|
||||
simp = self.simplices[sidx]
|
||||
verts = self.vertices[simp]
|
||||
|
||||
self.svolumes[sidx] = tet_volume(verts[0], verts[1], verts[2], verts[3])
|
||||
|
||||
for v in simp:
|
||||
self.pvolumes[v] += self.svolumes[sidx] / 4.0
|
||||
|
||||
|
||||
if verbose:
|
||||
print(' Done!', end ='')
|
||||
mtimer.end()
|
||||
|
||||
return self.pvolumes
|
||||
|
||||
def need_adjacentfaces(self, verbose=False):
|
||||
|
||||
'''
|
||||
Find adjacent faces for each vertex
|
||||
as list of lists
|
||||
'''
|
||||
|
||||
if len(self.adjacent_faces) != 0:
|
||||
return self.adjacent_faces
|
||||
|
||||
if verbose:
|
||||
print(' Computing adjacent_faces...', end ='')
|
||||
sys.stdout.flush()
|
||||
mtimer = Timer()
|
||||
|
||||
numadjacentfaces = numpy.zeros(self.nvertices, dtype=int)
|
||||
for f in self.simplices:
|
||||
for i in xrange(3):
|
||||
numadjacentfaces[f[i]] += 1
|
||||
|
||||
# can be optimized further by avoiding "append"?
|
||||
self.adjacent_faces = [[] for _ in xrange(self.nvertices)]
|
||||
for fidx in xrange(self.nsimplices):
|
||||
for i in xrange(3):
|
||||
self.adjacent_faces[self.simplices[fidx, i]].append(fidx)
|
||||
|
||||
if verbose:
|
||||
print('Done!', end ='')
|
||||
mtimer.end()
|
||||
|
||||
return self.adjacent_faces
|
||||
|
||||
def need_acrossedge(self, verbose=False):
|
||||
|
||||
'''
|
||||
Find adjacent faces for each face (across each edge)
|
||||
as ndarray of ints: shape (nsimplex, 3)
|
||||
-1 denotes a face on the boundary (no face across edge)
|
||||
'''
|
||||
if self.across_edge.shape != (0,0):
|
||||
return self.across_edge
|
||||
|
||||
self.need_adjacentfaces(verbose)
|
||||
|
||||
if verbose:
|
||||
print(' Computing across_edge...', end ='')
|
||||
sys.stdout.flush()
|
||||
mtimer = Timer()
|
||||
|
||||
self.across_edge = -1 * numpy.ones((self.nsimplices, 3), dtype=numpy.int)
|
||||
|
||||
for fidx in xrange(self.nsimplices):
|
||||
for i in xrange(3):
|
||||
|
||||
if self.across_edge[fidx, i] != -1:
|
||||
continue
|
||||
|
||||
v1 = self.simplices[fidx, (i+1)%3]
|
||||
v2 = self.simplices[fidx, (i+2)%3]
|
||||
|
||||
for other in self.adjacent_faces[v1]:
|
||||
|
||||
if other == fidx:
|
||||
continue
|
||||
|
||||
if other not in self.adjacent_faces[v2]:
|
||||
continue
|
||||
|
||||
oface = self.simplices[other]
|
||||
j = numpy.where(oface == v1)[0]
|
||||
j = (j+1)%3
|
||||
|
||||
if oface[(j+1)%3] != v2:
|
||||
continue
|
||||
|
||||
self.across_edge[fidx, i] = other
|
||||
self.across_edge[other, j] = fidx
|
||||
|
||||
if verbose:
|
||||
print(' Done!', end ='')
|
||||
mtimer.end()
|
||||
return self.across_edge
|
||||
|
||||
def need_boundary(self, verbose=False):
|
||||
|
||||
'''
|
||||
Find boundary of the triangulation
|
||||
as collection of boundary edeges as an array of [face, k]
|
||||
'''
|
||||
|
||||
if self.bedges.shape != (0,0):
|
||||
return self.bedges
|
||||
|
||||
self.need_acrossedge(verbose)
|
||||
|
||||
if verbose:
|
||||
print(' Computing the boundary...', end ='')
|
||||
sys.stdout.flush()
|
||||
mtimer = Timer()
|
||||
|
||||
# find all boundary faces and edges
|
||||
bfaces = [fidx for fidx in xrange(self.nsimplices) if -1 in self.across_edge[fidx]]
|
||||
bedges = []
|
||||
|
||||
for fidx in bfaces:
|
||||
face = self.simplices[fidx]
|
||||
nbrs = self.across_edge[fidx]
|
||||
bedges.extend( [[fidx, k] for k in xrange(3) if nbrs[k] == -1] )
|
||||
|
||||
self.bedges = numpy.array(bedges)
|
||||
if verbose:
|
||||
print(' Done! found', self.bedges.shape[0], 'boundary edges', end ='')
|
||||
mtimer.end()
|
||||
|
||||
return self.bedges
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# --------------------------------------------------------------------------
|
||||
def __init__(self, **kwargs):
|
||||
|
||||
'''
|
||||
kwargs:
|
||||
vertices: ndarray of shape (nverts, dim) # dim = 2,3
|
||||
simplices: ndarray of shape (nfaces, dim+1)
|
||||
verbose: verbosity level
|
||||
'''
|
||||
|
||||
args = kwargs.keys()
|
||||
|
||||
if 'vertices' not in args:
|
||||
raise SyntaxError("Mesh object needs vertex data")
|
||||
|
||||
verbose = 2
|
||||
if 'verbose' in args:
|
||||
verbose = kwargs['verbose']
|
||||
|
||||
self.vertices = kwargs['vertices']
|
||||
self.dim = self.vertices.shape[1]
|
||||
self.nvertices = self.vertices.shape[0]
|
||||
|
||||
if self.dim != 2 and self.dim != 3:
|
||||
raise SyntaxError("Mesh object works for 2D and 3D only")
|
||||
|
||||
if verbose > 0:
|
||||
print(' Initializing', self.dim, 'D mesh with', self.nvertices, 'vertices...', end ='')
|
||||
sys.stdout.flush()
|
||||
mtimer = Timer()
|
||||
|
||||
if verbose > 1:
|
||||
print('')
|
||||
|
||||
# create simplices if needed
|
||||
if 'simplices' in args:
|
||||
self.Delaunay = None
|
||||
self.simplices = kwargs['simplices']
|
||||
if verbose > 1:
|
||||
print(' got', self.simplices.shape[0], 'simplices')
|
||||
else:
|
||||
if verbose > 1:
|
||||
print(' creating Delaunay mesh...')
|
||||
sys.stdout.flush()
|
||||
|
||||
self.Delaunay = spatial.Delaunay(self.vertices)
|
||||
self.simplices = self.Delaunay.simplices
|
||||
|
||||
if verbose > 1:
|
||||
print(' Done! created', self.simplices.shape[0], 'simplices')
|
||||
sys.stdout.flush()
|
||||
|
||||
self.nsimplices = self.simplices.shape[0]
|
||||
if self.dim != self.simplices.shape[1]-1:
|
||||
pdim = self.dim
|
||||
sdim = self.simplices.shape[1]-1
|
||||
raise SyntaxError("Dimension mismatch! pdim = "+str(pdim)+" and sdim = "+str(sdim)+" do not match!")
|
||||
|
||||
self.adjacent_faces = []
|
||||
self.across_edge = numpy.empty((0,0))
|
||||
self.bedges = numpy.empty((0,0))
|
||||
self.pvolumes = numpy.empty(0)
|
||||
|
||||
self.need_volumes(verbose > 1)
|
||||
self.need_boundary(verbose > 1)
|
||||
#self.need_meshmatrices(verbose > 1)
|
||||
|
||||
if verbose > 0:
|
||||
print(' Done!', end ='')
|
||||
mtimer.end()
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# --------------------------------------------------------------------------
|
||||
def divcurl(self, vfield, verbose=False):
|
||||
|
||||
if vfield.shape != (self.nsimplices, self.dim):
|
||||
print('UnstructuredGrid.divcurl', vfield.shape, self.dim, self.vertices.shape, self.simplices.shape)
|
||||
raise ValueError("UnstructuredGrid requires a valid-dimensional vector field")
|
||||
|
||||
if verbose:
|
||||
print(' Computing divcurl...', end ='')
|
||||
sys.stdout.flush()
|
||||
mtimer = Timer()
|
||||
|
||||
div = numpy.zeros(self.nvertices)
|
||||
curlw = numpy.zeros(self.nvertices)
|
||||
|
||||
# for each face
|
||||
for sidx in xrange(self.nsimplices):
|
||||
|
||||
simp = self.simplices[sidx]
|
||||
|
||||
for k in xrange(3):
|
||||
|
||||
v = simp[k]
|
||||
a = simp[ (k+1)%3 ]
|
||||
b = simp[ (k+2)%3 ]
|
||||
|
||||
# normal and tangent vectors
|
||||
tvec = self.vertices[b] - self.vertices[a] # counterclockwise
|
||||
nvec = numpy.array([-tvec[1], tvec[0]]) # inward
|
||||
|
||||
dn = numpy.dot(nvec, vfield[sidx])
|
||||
tn = numpy.dot(tvec, vfield[sidx])
|
||||
|
||||
div[v] += dn
|
||||
curlw[v] += tn
|
||||
|
||||
# fix for boundary edges
|
||||
for bedge in self.bedges:
|
||||
|
||||
sidx = bedge[0]
|
||||
eidx = bedge[1]
|
||||
|
||||
a = self.simplices[sidx][(eidx+1)%3]
|
||||
b = self.simplices[sidx][(eidx+2)%3]
|
||||
|
||||
tvec = self.vertices[b] - self.vertices[a] # counterclockwise
|
||||
nvec = numpy.array([-tvec[1], tvec[0]]) # inward
|
||||
|
||||
dn = numpy.dot(nvec, vfield[sidx])
|
||||
dt = numpy.dot(tvec, vfield[sidx])
|
||||
|
||||
div[a] += dn
|
||||
div[b] += dn
|
||||
curlw[a] += tn
|
||||
curlw[b] += tn
|
||||
|
||||
div = -0.5 * div
|
||||
curlw = 0.5 * curlw
|
||||
|
||||
if verbose:
|
||||
print(' Done!', end ='')
|
||||
mtimer.end()
|
||||
#print('\tdiv =', div.shape, div.min(), div.max())
|
||||
#print('\tcurlw =', curlw.shape, curlw.min(), curlw.max())
|
||||
|
||||
return (div, curlw)
|
||||
|
||||
def gradient(self, sfield, verbose=False):
|
||||
|
||||
if sfield.shape[0] != self.nvertices:
|
||||
print('UnstructuredGrid.gradient', self.dim, sfield.shape, self.vertices.shape, self.simplices.shape)
|
||||
raise ValueError("UnstructuredGrid requires a valid-dimensional vector field")
|
||||
|
||||
if verbose:
|
||||
print(' Computing gradient...', end ='')
|
||||
sys.stdout.flush()
|
||||
mtimer = Timer()
|
||||
|
||||
grad = numpy.zeros((self.nsimplices, self.dim))
|
||||
|
||||
# for 2D
|
||||
for sidx in xrange(self.nsimplices):
|
||||
|
||||
deb = sidx==0
|
||||
|
||||
simp = self.simplices[sidx]
|
||||
f = 0.5 / self.svolumes[sidx]
|
||||
|
||||
if deb:
|
||||
print(sidx, simp, self.svolumes[sidx], f)
|
||||
|
||||
for k in xrange(3):
|
||||
|
||||
v = simp[k]
|
||||
a = simp[ (k+1)%3 ]
|
||||
b = simp[ (k+2)%3 ]
|
||||
|
||||
# normal and tangent vectors
|
||||
tvec = self.vertices[b] - self.vertices[a] # counterclockwise
|
||||
nvec = numpy.array([-tvec[1], tvec[0]]) # inward
|
||||
|
||||
grad[sidx] += f * sfield[v] * nvec
|
||||
|
||||
if deb:
|
||||
print('\t', k, nvec, sfield[v],grad[sidx])
|
||||
|
||||
if verbose:
|
||||
print(' Done!', end ='')
|
||||
mtimer.end()
|
||||
return grad
|
||||
|
||||
def rotated_gradient(self, sfield):
|
||||
|
||||
rgrad = self.gradient(sfield)
|
||||
|
||||
rgrad[:,[0, 1]] = rgrad[:,[1, 0]]
|
||||
rgrad[:,0] *= -1.0
|
||||
|
||||
return rgrad
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# --------------------------------------------------------------------------
|
||||
@@ -1,7 +1,57 @@
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
CONFIG_DIR = Path(__file__).parent
|
||||
def _get_config_dir():
|
||||
"""Resolve config directory, robustly handling Nuitka compiled paths."""
|
||||
|
||||
# 1. Start with the directory of this file
|
||||
current_path = Path(__file__).parent
|
||||
|
||||
# 2. Check strict relative location (standard case)
|
||||
# If we are in src/orisys/configs/__init__.py, configs is current_path
|
||||
if (current_path / "prototype.json").exists():
|
||||
return current_path
|
||||
|
||||
# 3. Walk up the directory tree to find a 'configs' folder containing 'prototype.json'
|
||||
# This handles nested virtual paths like .../orisys/orisys/configs
|
||||
# We look for a physical directory that matches our expectation
|
||||
|
||||
search_path = current_path
|
||||
# Limit depth to avoid infinite loops, though unlikely with .parent
|
||||
for _ in range(5):
|
||||
# Check if 'configs' exists in this directory
|
||||
candidate = search_path / "configs"
|
||||
if candidate.exists() and candidate.is_dir():
|
||||
if (candidate / "prototype.json").exists():
|
||||
return candidate
|
||||
|
||||
# Also check if the current directory ITSELF is the configs dir we missed
|
||||
# (e.g. if we are somehow inside configs but missed check 1)
|
||||
if search_path.name == 'configs' and (search_path / "prototype.json").exists():
|
||||
return search_path
|
||||
|
||||
if search_path.parent == search_path: # Reached root
|
||||
break
|
||||
search_path = search_path.parent
|
||||
|
||||
# 4. Try finding 'orisys' package root via sys.modules (if loaded)
|
||||
try:
|
||||
if 'orisys' in sys.modules:
|
||||
mod = sys.modules['orisys']
|
||||
if hasattr(mod, '__file__') and mod.__file__:
|
||||
pkg_root = Path(mod.__file__).parent
|
||||
# pkg_root is .../site-packages/orisys
|
||||
candidate = pkg_root / "configs"
|
||||
if candidate.exists() and (candidate / "prototype.json").exists():
|
||||
return candidate
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Fallback to current path if nothing else works
|
||||
return current_path
|
||||
|
||||
CONFIG_DIR = _get_config_dir()
|
||||
|
||||
def load_config(config_name_or_path):
|
||||
"""
|
||||
|
||||
@@ -1,356 +0,0 @@
|
||||
"""
|
||||
Density estimation module for optical flow fields.
|
||||
|
||||
This module computes density maps from optical flow vectors using Gaussian kernel
|
||||
density estimation. It supports both CPU (NumPy) and CUDA (CuPy) implementations.
|
||||
"""
|
||||
import cupy as cp
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
# Global variable to store the baseline/zero level of density for invalid area correction
|
||||
zero_level = None
|
||||
|
||||
def getFilteredDensity(flow, use_cuda):
|
||||
"""
|
||||
Compute filtered density map from optical flow field.
|
||||
|
||||
This function computes a density map using Gaussian kernel density estimation
|
||||
and then applies filtering to correct invalid areas.
|
||||
|
||||
Args:
|
||||
flow: Optical flow field with shape (height, width, 2) containing flow vectors
|
||||
use_cuda: If True, use CUDA-accelerated computation (requires CuPy and CUDA)
|
||||
If False, use CPU-based NumPy implementation
|
||||
|
||||
Returns:
|
||||
np.ndarray: Filtered density map with shape (height, width), where negative
|
||||
values indicate regions of flow divergence
|
||||
"""
|
||||
# Gaussian kernel parameters for density estimation
|
||||
sigma = 3.0 # Standard deviation of Gaussian kernel
|
||||
r_max = 9.0 # Maximum radius for Gaussian contribution (cutoff distance)
|
||||
|
||||
# Compute density map using either CUDA or CPU implementation
|
||||
if use_cuda:
|
||||
density = getDensityCuda(flow, sigma, r_max)
|
||||
else:
|
||||
density = getDensity(flow, sigma, r_max)
|
||||
|
||||
# Apply filtering to correct invalid areas (regions with unreliable flow)
|
||||
density = correctInvalidArea(density)
|
||||
|
||||
return density
|
||||
|
||||
def correctInvalidArea(density):
|
||||
"""
|
||||
Correct invalid areas in the density map.
|
||||
|
||||
Invalid areas are regions where the density falls below a threshold (zero_level),
|
||||
which typically indicates unreliable or missing flow data. This function:
|
||||
1. Uses adaptive zero_level that updates gradually to handle changing conditions
|
||||
2. Identifies invalid regions using adaptive thresholding (below zero_level - margin)
|
||||
3. Fills invalid areas using morphological closing for smoother transitions
|
||||
4. Applies stronger edge-preserving smoothing to eliminate discontinuities
|
||||
|
||||
Args:
|
||||
density: Input density map with shape (height, width)
|
||||
|
||||
Returns:
|
||||
np.ndarray: Corrected density map with invalid areas filled and smoothed,
|
||||
with improved continuity to reduce black dots/discontinuities
|
||||
"""
|
||||
global zero_level
|
||||
|
||||
# Compute adaptive zero level from current frame's central region
|
||||
# This adapts to changing density values when sensor is pressed
|
||||
current_zero_level = np.mean(density[60:-60, 60:-60])
|
||||
|
||||
# Initialize or update zero_level with exponential moving average for stability
|
||||
if zero_level is None:
|
||||
zero_level = current_zero_level
|
||||
else:
|
||||
# Use exponential moving average: 90% old, 10% new (adapts slowly)
|
||||
zero_level = 0.9 * zero_level + 0.1 * current_zero_level
|
||||
|
||||
# Use a threshold slightly below zero level to avoid over-correction
|
||||
# This prevents valid but slightly lower density regions from being marked invalid
|
||||
threshold = zero_level - np.std(density[60:-60, 60:-60]) * 0.5
|
||||
invalid_ind = twoClustering(density, threshold)
|
||||
|
||||
# Fill invalid areas using morphological closing to smooth transitions
|
||||
# This fills small holes and provides better continuity
|
||||
if np.any(invalid_ind > 0):
|
||||
# Create a temporary copy with invalid areas filled with zero_level
|
||||
density_filled = density.copy()
|
||||
density_filled[np.where(invalid_ind > 0)] = zero_level
|
||||
|
||||
# Apply morphological closing to fill small holes and smooth boundaries
|
||||
# This helps connect nearby valid regions and eliminate black dots
|
||||
kernel = np.ones((5, 5), np.float32)
|
||||
density_closed = cv2.morphologyEx(density_filled, cv2.MORPH_CLOSE, kernel, iterations=1)
|
||||
|
||||
# Use closed version only for invalid regions, preserve original valid values
|
||||
density[np.where(invalid_ind > 0)] = density_closed[np.where(invalid_ind > 0)]
|
||||
|
||||
# Apply edge-preserving smoothing using guided filter with stronger smoothing
|
||||
# Increased radius and epsilon for better continuity and to eliminate discontinuities
|
||||
try:
|
||||
density = cv2.ximgproc.guidedFilter(density, density, 15, 0.01)
|
||||
except AttributeError:
|
||||
# Fallback to stronger Gaussian blur if ximgproc is not available
|
||||
# Use larger kernel for better smoothing of discontinuities
|
||||
density = cv2.GaussianBlur(density, (21, 21), 0)
|
||||
|
||||
return density
|
||||
|
||||
|
||||
def twoClustering(density, threshold):
|
||||
"""
|
||||
Perform binary clustering/segmentation of density map based on threshold.
|
||||
|
||||
This function creates a binary mask where pixels below the threshold are
|
||||
marked as invalid (value 1) and pixels above threshold are valid (value 0).
|
||||
|
||||
Args:
|
||||
density: Input density map
|
||||
threshold: Threshold value for segmentation
|
||||
|
||||
Returns:
|
||||
np.ndarray: Binary mask with dtype uint8, where 1 indicates invalid areas
|
||||
(density < threshold) and 0 indicates valid areas
|
||||
"""
|
||||
cluster = np.zeros_like(density, dtype='uint8')
|
||||
cluster[np.where(density < threshold)] = 1
|
||||
return cluster
|
||||
|
||||
|
||||
|
||||
def getDensity(flow, sigma, r_max):
|
||||
"""
|
||||
Compute Gaussian density map from optical flow vectors.
|
||||
Pure NumPy implementation without freud dependency.
|
||||
Vectorized for better performance.
|
||||
|
||||
Args:
|
||||
flow: Optical flow field with shape (height, width, 2)
|
||||
sigma: Gaussian kernel standard deviation
|
||||
r_max: Maximum radius for Gaussian contribution
|
||||
|
||||
Returns:
|
||||
Negative density map with same shape as flow field (height, width)
|
||||
"""
|
||||
# Pre-compute constants
|
||||
two_sigma2 = 2.0 * sigma * sigma
|
||||
gaussian_norm = 1.0 / (two_sigma2 * np.pi)
|
||||
|
||||
# Get flow dimensions
|
||||
height, width = flow.shape[:2]
|
||||
flow_x = flow[:, :, 0].astype(np.float32)
|
||||
flow_y = flow[:, :, 1].astype(np.float32)
|
||||
|
||||
# Create coordinate grids
|
||||
x_coords = np.arange(width, dtype=np.float32)
|
||||
y_coords = np.arange(height, dtype=np.float32)
|
||||
grid_x, grid_y = np.meshgrid(x_coords, y_coords)
|
||||
|
||||
# Compute destination positions after flow displacement
|
||||
dest_x = grid_x + flow_x
|
||||
dest_y = grid_y + flow_y
|
||||
|
||||
# Initialize result array
|
||||
result = np.zeros((height, width), dtype=np.float32)
|
||||
|
||||
# Compute density by iterating over each source pixel
|
||||
# For each pixel, apply its flow vector to get destination position,
|
||||
# then accumulate Gaussian contributions to nearby grid points
|
||||
for y in range(height):
|
||||
for x in range(width):
|
||||
# Get the destination position after applying flow displacement
|
||||
dest_x_val = dest_x[y, x]
|
||||
dest_y_val = dest_y[y, x]
|
||||
|
||||
# Determine the bounding box of grid points affected by this flow vector
|
||||
# These are grid points within r_max radius of the destination
|
||||
min_x = max(0, int(np.floor(dest_x_val - r_max)))
|
||||
max_x = min(width, int(np.ceil(dest_x_val + r_max)) + 1)
|
||||
min_y = max(0, int(np.floor(dest_y_val - r_max)))
|
||||
max_y = min(height, int(np.ceil(dest_y_val + r_max)) + 1)
|
||||
|
||||
# Skip if bounding box is empty or out of bounds
|
||||
if min_x >= max_x or min_y >= max_y:
|
||||
continue
|
||||
|
||||
# Create local coordinate grid for the affected region
|
||||
local_x = np.arange(min_x, max_x, dtype=np.float32)
|
||||
local_y = np.arange(min_y, max_y, dtype=np.float32)
|
||||
local_grid_x, local_grid_y = np.meshgrid(local_x, local_y)
|
||||
|
||||
# Compute squared distances from destination to each grid point
|
||||
dx = local_grid_x - dest_x_val
|
||||
dy = local_grid_y - dest_y_val
|
||||
distance_sq = dx * dx + dy * dy
|
||||
|
||||
# Create mask for points within r_max radius (circular neighborhood)
|
||||
mask = distance_sq <= (r_max * r_max)
|
||||
|
||||
# Compute Gaussian kernel values: exp(-d²/(2σ²)) / (2πσ²)
|
||||
# Higher values for closer points, decreasing with distance
|
||||
gaussian_vals = np.exp(-distance_sq / two_sigma2) * gaussian_norm
|
||||
gaussian_vals[~mask] = 0 # Zero out contributions beyond r_max
|
||||
|
||||
# Accumulate negative density (negative indicates flow divergence)
|
||||
# Each flow vector contributes to density at its destination region
|
||||
result[min_y:max_y, min_x:max_x] -= gaussian_vals
|
||||
|
||||
return result
|
||||
|
||||
|
||||
class CudaArrayInterface:
|
||||
"""
|
||||
Interface class to convert OpenCV CUDA GpuMat to CuPy-compatible array interface.
|
||||
|
||||
This allows seamless conversion between OpenCV's CUDA memory objects and CuPy arrays
|
||||
without copying data, by implementing the CUDA array interface protocol.
|
||||
"""
|
||||
def __init__(self, gpu_mat):
|
||||
"""
|
||||
Initialize CUDA array interface from OpenCV GpuMat.
|
||||
|
||||
Args:
|
||||
gpu_mat: OpenCV CUDA GpuMat object
|
||||
"""
|
||||
w, h = gpu_mat.size()
|
||||
|
||||
# Map OpenCV data types to NumPy/CuPy type strings
|
||||
type_map = {
|
||||
cv2.CV_8U: "u1", # unsigned 8-bit
|
||||
cv2.CV_8S: "i1", # signed 8-bit
|
||||
cv2.CV_16U: "u2", # unsigned 16-bit
|
||||
cv2.CV_16S: "i2", # signed 16-bit
|
||||
cv2.CV_32S: "i4", # signed 32-bit integer
|
||||
cv2.CV_32F: "f4", # 32-bit float
|
||||
cv2.CV_64F: "f8", # 64-bit float
|
||||
}
|
||||
|
||||
# Implement CUDA array interface protocol for zero-copy conversion
|
||||
self.__cuda_array_interface__ = {
|
||||
"version": 2,
|
||||
"shape": (h, w),
|
||||
"data": (gpu_mat.cudaPtr(), False), # CUDA memory pointer, not writable
|
||||
"typestr": type_map[gpu_mat.type()],
|
||||
"strides": (gpu_mat.step, gpu_mat.elemSize()), # Row stride and element size
|
||||
}
|
||||
|
||||
|
||||
# CUDA kernel code for accelerated density computation
|
||||
# This kernel computes Gaussian density contributions in parallel on GPU
|
||||
func_string = r"""
|
||||
extern "C"{
|
||||
// Compute 2D Gaussian kernel value
|
||||
// Args:
|
||||
// twoSigma2: 2 * sigma^2 (precomputed for efficiency)
|
||||
// dx, dy: Distance components from center
|
||||
// Returns: Gaussian kernel value normalized by (2*pi*sigma^2)
|
||||
__device__ float gaussian(float twoSigma2, float dx, float dy)
|
||||
{
|
||||
float pi = 3.1415926535897932384626;
|
||||
float distance = dx*dx+dy*dy; // Squared distance
|
||||
float result = expf(-distance/twoSigma2); // Gaussian exponential term
|
||||
result = result/(twoSigma2*pi); // Normalization factor
|
||||
return result;
|
||||
}
|
||||
|
||||
// CUDA kernel to compute density map from optical flow
|
||||
// Each thread processes one pixel in the flow field
|
||||
// Args:
|
||||
// flowX, flowY: Input flow field components (flattened)
|
||||
// result: Output density map (flattened, initialized to zero)
|
||||
// x_size, y_size: Dimensions of the flow field
|
||||
// twoSigma2: Precomputed 2 * sigma^2
|
||||
// r_max: Maximum radius for Gaussian contribution
|
||||
__global__ void density(const float *flowX, const float *flowY, float *result,
|
||||
int x_size, int y_size, float twoSigma2, float r_max)
|
||||
{
|
||||
// Get thread coordinates (each thread handles one source pixel)
|
||||
int x = threadIdx.x; // Column index
|
||||
int y = blockIdx.x; // Row index
|
||||
|
||||
// Compute destination position after applying flow displacement
|
||||
float dest_x = float(x) + flowX[y*x_size+x];
|
||||
float dest_y = float(y) + flowY[y*x_size+x];
|
||||
|
||||
// Iterate over nearby grid points within r_max radius
|
||||
for(int ind_x = ceilf(dest_x - r_max); ind_x < ceilf(dest_x + r_max); ++ind_x)
|
||||
for(int ind_y = ceilf(dest_y - r_max); ind_y < ceilf(dest_y + r_max); ++ind_y)
|
||||
{
|
||||
// Bounds checking
|
||||
if(ind_x >= x_size || ind_x < 0 || ind_y >= y_size || ind_y < 0)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
|
||||
// Compute distance from destination to grid point
|
||||
float dx = float(ind_x) - dest_x;
|
||||
float dy = float(ind_y) - dest_y;
|
||||
|
||||
// Accumulate Gaussian contribution (negative for divergence representation)
|
||||
atomicAdd(result + ind_y*x_size + ind_x, -gaussian(twoSigma2, dx, dy));
|
||||
}
|
||||
}
|
||||
}
|
||||
"""
|
||||
# Compile CUDA kernel into CuPy RawModule for execution
|
||||
funcupy = cp.RawModule(code=func_string).get_function("density")
|
||||
|
||||
|
||||
def getDensityCuda(flow, sigma, r_max):
|
||||
"""
|
||||
Compute Gaussian density map from optical flow using CUDA acceleration.
|
||||
|
||||
This function uses a custom CUDA kernel for parallel computation of density maps,
|
||||
providing significant speedup over CPU implementation for large flow fields.
|
||||
|
||||
Args:
|
||||
flow: Optical flow field with shape (height, width, 2), can be numpy array
|
||||
or OpenCV CUDA GpuMat
|
||||
sigma: Gaussian kernel standard deviation
|
||||
r_max: Maximum radius for Gaussian contribution (cutoff distance)
|
||||
|
||||
Returns:
|
||||
np.ndarray: Density map with shape (height, width), transferred from GPU to CPU
|
||||
|
||||
Note:
|
||||
Requires CuPy and CUDA-capable GPU. Falls back to CPU if unavailable.
|
||||
"""
|
||||
# Precompute 2*sigma^2 for efficiency
|
||||
twoSigma2 = cp.float32(2.0 * sigma**2)
|
||||
r_max = cp.float32(r_max)
|
||||
|
||||
# Convert input to CuPy arrays
|
||||
# Handle both numpy arrays and OpenCV CUDA GpuMat
|
||||
if type(flow).__module__ == np.__name__:
|
||||
# NumPy array: convert directly to CuPy
|
||||
flowx = cp.array(flow[:, :, 0], dtype=cp.float32)
|
||||
flowy = cp.array(flow[:, :, 1], dtype=cp.float32)
|
||||
else:
|
||||
# OpenCV CUDA GpuMat: use zero-copy interface
|
||||
flowx, flowy = cv2.cuda.split(flow)
|
||||
flowx = cp.array(CudaArrayInterface(flowx), dtype=cp.float32, copy=True)
|
||||
flowy = cp.array(CudaArrayInterface(flowy), dtype=cp.float32, copy=True)
|
||||
|
||||
# Initialize output array on GPU
|
||||
result = cp.zeros((flowx.shape[0], flowx.shape[1]), dtype=cp.float32)
|
||||
|
||||
# Configure CUDA kernel launch parameters
|
||||
block_dim = (flowx.shape[1], ) # Threads per block (one thread per column)
|
||||
grid = (flowx.shape[0], ) # Number of blocks (one block per row)
|
||||
|
||||
# Launch CUDA kernel
|
||||
funcupy(grid, block_dim, (flowx, flowy, result,
|
||||
cp.int32(flowx.shape[1]),
|
||||
cp.int32(flowx.shape[0]),
|
||||
twoSigma2, r_max))
|
||||
|
||||
# Transfer result from GPU to CPU
|
||||
return result.get()
|
||||
@@ -1,798 +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
|
||||
from .nPoisson import PoissonSolver
|
||||
from .StructuredGrid import StructuredGrid
|
||||
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 decompose_batch(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 using GPU
|
||||
# (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 using batch GPU solver
|
||||
# batch_inputs = [self.div, self.curlw]
|
||||
# batch_results = self.psolver.solve_gpu_batch(batch_inputs, verbose > 0)
|
||||
|
||||
# self.nD = batch_results[0]
|
||||
# self.nRu = batch_results[1]
|
||||
|
||||
# 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 using batch GPU solver
|
||||
# batch_inputs = [self.div, self.curlu, self.curlv, self.curlw]
|
||||
# batch_results = self.psolver.solve_gpu_batch(batch_inputs, verbose > 0)
|
||||
|
||||
# self.nD = batch_results[0]
|
||||
# self.nRu = batch_results[1] * -1.0
|
||||
# self.nRv = batch_results[2] * -1.0
|
||||
# self.nRw = batch_results[3] * -1.0
|
||||
|
||||
# if verbose > 0:
|
||||
# print( 'nD =', self.nD.shape, self.nD.min(), self.nD.max())
|
||||
# print( 'nRu =', self.nRu.shape, self.nRu.min(), self.nRu.max())
|
||||
# print( 'nRv =', self.nRv.shape, self.nRv.min(), self.nRv.max())
|
||||
# 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 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(self, vfield, verbose=False):
|
||||
# """
|
||||
# GPU-accelerated force computation using CuPy.
|
||||
# Falls back to CPU if CuPy is not available.
|
||||
|
||||
# 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 (uses GPU internally)
|
||||
# self.decompose(vfield, verbose)
|
||||
|
||||
# # Initialize cache if needed or shape changed
|
||||
# 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:
|
||||
# self._init_force_gpu_cache(shape)
|
||||
|
||||
# # Transfer data to GPU (consider keeping on GPU in decompose for further optimization)
|
||||
# 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)
|
||||
|
||||
# # Compute all sums in parallel - fused for better performance
|
||||
# # Sn: sum of absolute values of divergent field (both components)
|
||||
# Sn_gpu = cp.sum(cp.abs(d_gpu))
|
||||
|
||||
# # Stx and Sty: sum of harmonic field components
|
||||
# Stx_gpu = cp.sum(h_gpu[:, :, 0])
|
||||
# Sty_gpu = cp.sum(h_gpu[:, :, 1])
|
||||
|
||||
# # Find max/min locations of nRu using single pass
|
||||
# nR_max_idx = int(cp.argmax(nRu_gpu))
|
||||
# nR_min_idx = int(cp.argmin(nRu_gpu))
|
||||
|
||||
# # 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_min_loc = numpy.unravel_index(nR_min_idx, nRu_gpu.shape)
|
||||
|
||||
# # Use cached grid arrays
|
||||
# mgrid_0 = self._force_gpu_cache['mgrid_0_bc']
|
||||
# mgrid_1 = self._force_gpu_cache['mgrid_1_bc']
|
||||
|
||||
# # Compute torque using fused operations
|
||||
# # 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])
|
||||
# min_0, min_1 = float(nR_min_loc[0]), float(nR_min_loc[1])
|
||||
|
||||
# r_x = r_gpu[:, :, 0]
|
||||
# r_y = r_gpu[:, :, 1]
|
||||
|
||||
# # Fused cross product computation
|
||||
# # 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
|
||||
# Stor_gpu = 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
|
||||
# )
|
||||
|
||||
# # Single batch transfer to CPU (more efficient than 4 separate transfers)
|
||||
# results = cp.asnumpy(cp.array([Sn_gpu, Stx_gpu, Sty_gpu, Stor_gpu]))
|
||||
|
||||
# self.Sn = float(results[0])
|
||||
# self.Stx = float(results[1])
|
||||
# self.Sty = float(results[2])
|
||||
# self.Stor = float(results[3])
|
||||
|
||||
# output = results
|
||||
|
||||
# # Check thresholds
|
||||
# flag = (self.Sn >= 300 and self.Stx >= 300 and self.Sty >= 300 and numpy.abs(self.Stor) >= 1000)
|
||||
|
||||
# return flag, output
|
||||
|
||||
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_force_gpu_stream(self, vfield, verbose=False):
|
||||
# """
|
||||
# Fully optimized GPU force computation using:
|
||||
# - decompose_gpu to keep data on GPU throughout
|
||||
# - CUDA streams for parallel reductions
|
||||
# - Fused kernel operations
|
||||
|
||||
# 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)
|
||||
|
||||
# # Use GPU decomposition to keep data on GPU
|
||||
# gpu_result = self.decompose_gpu(vfield, verbose)
|
||||
|
||||
# if gpu_result is None:
|
||||
# # Fallback for 3D or if decompose_gpu failed
|
||||
# return self.get_force_gpu(vfield, verbose)
|
||||
|
||||
# d_gpu, h_gpu, r_gpu, nRu_gpu = gpu_result
|
||||
|
||||
# # Initialize cache if needed or shape changed
|
||||
# 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:
|
||||
# self._init_force_gpu_cache(shape)
|
||||
|
||||
# # Use CUDA streams for parallel computation if available
|
||||
# if STREAMS_AVAILABLE:
|
||||
# stream1 = Stream(non_blocking=True)
|
||||
# stream2 = Stream(non_blocking=True)
|
||||
# stream3 = Stream(non_blocking=True)
|
||||
|
||||
# # Stream 1: Compute Sn
|
||||
# with stream1:
|
||||
# Sn_gpu = cp.sum(cp.abs(d_gpu))
|
||||
|
||||
# # Stream 2: Compute Stx, Sty
|
||||
# with stream2:
|
||||
# Stx_gpu = cp.sum(h_gpu[:, :, 0])
|
||||
# Sty_gpu = cp.sum(h_gpu[:, :, 1])
|
||||
|
||||
# # Stream 3: Find extrema and compute torque
|
||||
# with stream3:
|
||||
# nR_max_idx = int(cp.argmax(nRu_gpu))
|
||||
# nR_min_idx = int(cp.argmin(nRu_gpu))
|
||||
|
||||
# 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]
|
||||
|
||||
# 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]
|
||||
|
||||
# Stor_gpu = 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
|
||||
# )
|
||||
|
||||
# # Synchronize all streams
|
||||
# stream1.synchronize()
|
||||
# stream2.synchronize()
|
||||
# stream3.synchronize()
|
||||
# else:
|
||||
# # Without streams, compute sequentially but still on GPU
|
||||
# Sn_gpu = cp.sum(cp.abs(d_gpu))
|
||||
# Stx_gpu = cp.sum(h_gpu[:, :, 0])
|
||||
# Sty_gpu = cp.sum(h_gpu[:, :, 1])
|
||||
|
||||
# nR_max_idx = int(cp.argmax(nRu_gpu))
|
||||
# nR_min_idx = int(cp.argmin(nRu_gpu))
|
||||
|
||||
# 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]
|
||||
|
||||
# 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]
|
||||
|
||||
# Stor_gpu = 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
|
||||
# )
|
||||
|
||||
# # Single batch transfer to CPU
|
||||
# results = cp.asnumpy(cp.array([Sn_gpu, Stx_gpu, Sty_gpu, Stor_gpu]))
|
||||
|
||||
# self.Sn = float(results[0])
|
||||
# self.Stx = float(results[1])
|
||||
# self.Sty = float(results[2])
|
||||
# self.Stor = float(results[3])
|
||||
|
||||
# output = results
|
||||
|
||||
# # Check thresholds
|
||||
# flag = (self.Sn >= 300 and self.Stx >= 300 and self.Sty >= 300 and numpy.abs(self.Stor) >= 1000)
|
||||
|
||||
# return flag, output
|
||||
|
||||
# def get_force_gpu_kernel(self, vfield, verbose=False):
|
||||
# """
|
||||
# Maximum performance GPU force computation using:
|
||||
# - decompose_gpu to keep data on GPU throughout
|
||||
# - Custom CUDA kernel for fused torque computation
|
||||
# - Minimized memory allocations
|
||||
|
||||
# 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)
|
||||
|
||||
# # Use GPU decomposition to keep data on GPU
|
||||
# gpu_result = self.decompose_gpu(vfield, verbose)
|
||||
|
||||
# if gpu_result is None:
|
||||
# return self.get_force_gpu(vfield, verbose)
|
||||
|
||||
# d_gpu, h_gpu, r_gpu, nRu_gpu = gpu_result
|
||||
|
||||
# # Initialize cache if needed
|
||||
# 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:
|
||||
# self._init_force_gpu_cache(shape)
|
||||
|
||||
# # Compute Sn using L1 norm (faster than sum(abs()))
|
||||
# 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
|
||||
# nR_max_idx = int(cp.argmax(nRu_gpu))
|
||||
# nR_min_idx = int(cp.argmin(nRu_gpu))
|
||||
|
||||
# max_0 = 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_1 = cp.float32(nR_min_idx % shape[1])
|
||||
|
||||
# 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]
|
||||
|
||||
# # Use custom kernel if available for fused torque computation
|
||||
# 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)
|
||||
# Stor = float(cp.sum(torque_vals))
|
||||
# else:
|
||||
# # Fallback to standard computation
|
||||
# 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 get_shear(self):
|
||||
|
||||
return self.h
|
||||
|
||||
def get_normal(self):
|
||||
|
||||
return self.d
|
||||
|
||||
# def plot_2d_information(self,value):
|
||||
# fig,ax = plt.subplots()
|
||||
# x,y = [],[]
|
||||
# for i in range (1000000000000):
|
||||
# x.append(i)
|
||||
# y.append(value)
|
||||
# ax.cla()
|
||||
# ax.plot(x,y,'b',lw=1)
|
||||
|
||||
# plt.pause(0.1)
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -1,519 +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
|
||||
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
|
||||
|
||||
|
||||
'''
|
||||
# ------------------------------------------------------------------------------
|
||||
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 (self.dim > 3):
|
||||
raise ValueError("Poisson Solver works only for 1, 2, or 3 dimensions")
|
||||
|
||||
self.ready = False
|
||||
self._gpu_cache = None # Cache for GPU arrays
|
||||
# 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
|
||||
if self.stype == 'F':
|
||||
p = signal.fftconvolve(f, self.G, mode='same')
|
||||
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):
|
||||
"""
|
||||
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']
|
||||
|
||||
# Transfer result back to CPU
|
||||
p = cp.asnumpy(p_gpu)
|
||||
|
||||
elif self.stype == 'S':
|
||||
raise SyntaxError('Spatial solver not suggested for structured grids')
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# 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)
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
if verbose:
|
||||
print('Poisson solver (GPU) finished')
|
||||
gtimer.end()
|
||||
|
||||
return p
|
||||
|
||||
def solve_gpu_batch(self, f_list, verbose=False):
|
||||
"""
|
||||
Batch GPU-accelerated Poisson solver for multiple inputs.
|
||||
Solves multiple Poisson equations in a single GPU call.
|
||||
|
||||
Args:
|
||||
f_list: List of input fields (numpy arrays)
|
||||
verbose: Print timing information
|
||||
|
||||
Returns:
|
||||
List of solutions (numpy arrays on CPU)
|
||||
"""
|
||||
if not CUPY_AVAILABLE:
|
||||
if verbose:
|
||||
print("CuPy not available, falling back to CPU solver")
|
||||
return [self.solve(f, verbose) for f in f_list]
|
||||
|
||||
if not self.ready:
|
||||
self.prepare(verbose)
|
||||
|
||||
gtimer = Timer()
|
||||
if verbose:
|
||||
print(f'\nSolving {len(f_list)} Poisson Eq. on GPU (batched)')
|
||||
|
||||
# Verify all inputs have same shape
|
||||
fshape = f_list[0].shape
|
||||
for f in f_list:
|
||||
if f.shape != fshape:
|
||||
raise ValueError("All input fields must have the same shape for batch processing")
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# regular grid
|
||||
if self.ptype == 'G':
|
||||
if 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)
|
||||
}
|
||||
|
||||
results = []
|
||||
|
||||
# Process all inputs in batch
|
||||
for i, f in enumerate(f_list):
|
||||
# Transfer input to GPU
|
||||
f_gpu = cp.asarray(f, dtype=cp.float32)
|
||||
|
||||
# Perform FFT convolution on GPU
|
||||
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']
|
||||
|
||||
# Transfer result back to CPU
|
||||
results.append(cp.asnumpy(p_gpu))
|
||||
|
||||
if verbose:
|
||||
print(f' - Solved {i+1}/{len(f_list)}')
|
||||
|
||||
if verbose:
|
||||
print('Poisson solver (GPU batch) finished')
|
||||
gtimer.end()
|
||||
|
||||
return results
|
||||
|
||||
elif self.stype == 'S':
|
||||
raise SyntaxError('Spatial solver not suggested for structured grids')
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# 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) for f in f_list]
|
||||
|
||||
# If we reach here, it's an unsupported case
|
||||
raise ValueError(f"Unsupported point type '{self.ptype}' for GPU batch solver")
|
||||
|
||||
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()
|
||||
|
||||
|
||||
@@ -1,59 +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
|
||||
|
||||
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()
|
||||
@@ -1,6 +1,6 @@
|
||||
import cv2
|
||||
import numpy as np
|
||||
from .nHHD import *
|
||||
# from .HHD import *
|
||||
# import src.utils as utils
|
||||
|
||||
|
||||
@@ -172,10 +172,10 @@ class Tracker:
|
||||
self.__prev_fine_flow)
|
||||
return self.__prev_fine_flow
|
||||
|
||||
def track_2dto3d(self, flow):
|
||||
decompose_obj = nHHD(grid=(flow.shape[0], flow.shape[1]), spacings=(18, 18))
|
||||
flag,output = decompose_obj.get_force(flow, verbose=0)
|
||||
return output
|
||||
# def track_2dto3d(self, flow):
|
||||
# decompose_obj = nHHD(grid=(flow.shape[0], flow.shape[1]), spacings=(18, 18))
|
||||
# flag,output = decompose_obj.get_force(flow, verbose=0)
|
||||
# return output
|
||||
|
||||
def __getForwardBackwardError(self, I0, I1, flow):
|
||||
if self.cuda:
|
||||
|
||||
Reference in New Issue
Block a user