combine HHD

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 .nPoisson import PoissonSolver
-from .StructuredGrid import StructuredGrid
+# 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()
+    
+