diff --git a/src/orisys/HHD.py b/src/orisys/HHD.py index eb4be2b..be7f1c5 100644 --- a/src/orisys/HHD.py +++ b/src/orisys/HHD.py @@ -22,10 +22,11 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. ''' import numpy -from .nPoisson import PoissonSolver -from .StructuredGrid import StructuredGrid +# from .nPoisson import PoissonSolver +# from .StructuredGrid import StructuredGrid import timeit import sys +from scipy import signal, spatial # from .UnstructuredGrid import UnstructuredGrid # import matplotlib.pyplot as plt @@ -855,3 +856,501 @@ class StructuredGrid(object): # -------------------------------------------------------------------------- # -------------------------------------------------------------------------- + + +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() + + \ No newline at end of file