import numpy from scipy import spatial from ..backend_availability import CUPY_AVAILABLE from .config import MaskedPoissonConfig from .domain import MaskedDomain, normalize_mask from .solvers.capacitance import CapacitanceFFTSolver from .solvers.masked_sparse import MaskedSparseSolver from .solvers.rectangular import RectangularFFTSolver from .timing import Timer ''' # ------------------------------------------------------------------------------ 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: 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") self.dtype = numpy.float32 if "dtype" in args: self.dtype = kwargs["dtype"] if "grid" in args: self.ptype = "G" self.stype = "F" 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" 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 "poisson_backend" in args: self.poisson_backend = kwargs["poisson_backend"] if self.poisson_backend not in ("cpu", "opencl", "cupy"): raise ValueError("poisson_backend must be one of: 'cpu', 'opencl', 'cupy'") else: self.poisson_backend = "cpu" self.masked_config = MaskedPoissonConfig.from_kwargs(kwargs) self.masked_poisson_solver = self.masked_config.method self.masked_poisson_max_boundary = self.masked_config.max_boundary self.masked_poisson_boundary_stride = self.masked_config.boundary_stride self.masked_poisson_max_iters = self.masked_config.max_iters self.masked_poisson_tol = self.masked_config.tol self.masked_poisson_reg = self.masked_config.reg if self.dim > 3: raise ValueError("Poisson Solver works only for 1, 2, or 3 dimensions") self.ready = False self.G = None self._masked_cache = None self._masked_domain = None self._opencv_cache = None self._rect_solver = None self._sparse_solver = MaskedSparseSolver() self._cap_solver = None if self.ptype == "G": self._rect_solver = RectangularFFTSolver( grid=self.gdims, spacings=self.gdx, poisson_backend=self.poisson_backend, dtype=self.dtype, stype=self.stype, ) self._cap_solver = CapacitanceFFTSolver(self.masked_config, self._rect_solver) if self.poisson_backend == "cupy": self._gpu_cache = None if self.poisson_backend == "opencl": self._opencl_cache = None self._opencl_queue = None self._opencl_context = None self._opencl_last_timing = None def prepare(self, verbose=False): """Initialize Green's function / distance matrix for the domain. verbose: print timing and setup info returns: None (sets self.ready) """ gtimer = Timer() if verbose: print("\nInitializing Poisson solver, type =", self.stype) if self.ptype == "G" and self.stype == "F": self._rect_solver.prepare(verbose) self.G = self._rect_solver.G self.ready = self._rect_solver.ready elif self.stype == "F" and self.ptype == "P": raise SyntaxError("Frequncy solver not supported for unstructured points") elif self.stype == "S" and self.ptype == "G": raise SyntaxError("Spatial solver not suggested for structured grids") elif self.stype == "S" and self.ptype == "P": if self.dim == 1 and len(self.points.shape) == 1: self.points = self.points[:, numpy.newaxis] ltimer = Timer() self.G = spatial.distance.cdist(self.points, self.points) 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() 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) self.ready = True if verbose: print("Poisson solver initialized", end="") gtimer.end() def _normalize_mask(self, mask, shape): return normalize_mask(mask, shape) def _masked_cache_matches(self, mask): return self._masked_domain is not None and self._masked_domain.matches(mask) def _get_or_build_domain(self, mask): if self._masked_cache_matches(mask): return self._masked_domain if self.dim != 2: raise NotImplementedError("Masked Poisson solve is currently implemented for 2D grids only") domain = MaskedDomain.build(mask, self.gdx) capfft = None if self.masked_config.method == "capfft": capfft = self._cap_solver.prepare(domain) self._masked_domain = domain self._masked_cache = domain.to_cache_dict(capfft=capfft) return domain def _build_masked_cache(self, mask): """Legacy wrapper; returns dict cache for tests.""" self._get_or_build_domain(mask) return self._masked_cache def _compute_masked_residual(self, field, rhs_active, cache): """Legacy wrapper for capacitance/residual tests.""" domain = self._masked_domain if domain is None or not domain.matches(cache["mask"]): domain = MaskedDomain.build(cache["mask"], self.gdx) return domain.residual(field, rhs_active) def _solve_masked(self, f, mask, verbose=False): mask_arr = self._normalize_mask(mask, f.shape) if mask_arr is None: return None if not numpy.any(mask_arr): return numpy.zeros_like(f, dtype=numpy.float32) domain = self._get_or_build_domain(mask_arr) if self.masked_config.method == "capfft": return self._cap_solver.solve(f, domain) return self._sparse_solver.solve(f, domain) def solve_masked_pair(self, f0, f1, mask, verbose=False): """Solve two masked RHS fields, using a batched sparse path when available.""" mask_arr = self._normalize_mask(mask, f0.shape) if mask_arr is None: return None, None if f1.shape != f0.shape: raise ValueError("Both masked RHS fields must have the same shape") if not numpy.any(mask_arr): empty = numpy.zeros_like(f0, dtype=numpy.float32) return empty, empty.copy() domain = self._get_or_build_domain(mask_arr) if self.masked_config.method == "capfft": return self._cap_solver.solve(f0, domain), self._cap_solver.solve(f1, domain) return self._sparse_solver.solve_pair(f0, f1, domain) def solve(self, f, verbose=False, mask=None): """Solve Poisson equation on the prepared grid. f: (H, W) scalar source field (divergence or curl weight) mask: optional (H, W) bool; True = active UV/domain pixels returns: (H, W) scalar potential; zero outside mask when mask is set """ if not self.ready: self.prepare(verbose) fshape = f.shape mask_arr = self._normalize_mask(mask, fshape) gtimer = Timer() if verbose: print("\nSolving Poisson Eq.") if self.ptype == "G": if self.gdims != fshape: print(self.gdims, fshape) raise ValueError("Shape of function should match shape of grid") if mask_arr is not None: p = self._solve_masked(f, mask_arr, verbose) if verbose: print("Poisson solver finished") gtimer.end() return p if self.stype == "F": p = self._rect_solver.solve(f, verbose=False) else: raise SyntaxError("Spatial solver not suggested for structured grids") 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") if self.stype == "S": p = numpy.multiply(f, self.pvolumes) numpy.dot(self.G, p, p) else: raise SyntaxError("Frequncy solver not supported for unstructured points") if verbose: print("Poisson solver finished") gtimer.end() return p def solve_gpu(self, f, verbose=False, return_gpu=False): """GPU FFT Poisson solve (full rectangular grid only). f: (H, W) scalar source field return_gpu: return CuPy array if True returns: (H, W) scalar potential """ 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) if self.ptype != "G" or self.stype != "F": raise SyntaxError("GPU solver applies to structured FFT grids only") result = self._rect_solver.solve_gpu(f, verbose=verbose, return_gpu=return_gpu) self._gpu_cache = self._rect_solver._gpu_cache return result def _solve_gpu_internal(self, f_gpu): """Internal GPU solve; input and output are CuPy arrays.""" if not self.ready: self.prepare(False) return self._rect_solver._solve_gpu_internal(f_gpu) def clear_gpu_cache(self): """Release CuPy FFT cache buffers.""" if self._rect_solver is not None: self._rect_solver.clear_gpu_cache() self._gpu_cache = None def solve_opencl_batch2(self, f1, f2, verbose=False, profile=False): """Batched OpenCL FFT Poisson solve for two 2D fields. f1, f2: (H, W) scalar source fields returns: (p1, p2) each (H, W) scalar potential """ if not self.ready: self.prepare(verbose) p1, p2 = self._rect_solver.solve_opencl_batch2(f1, f2, verbose=verbose, profile=profile) self._opencl_cache = self._rect_solver._opencl_cache self._opencl_queue = self._rect_solver._opencl_queue self._opencl_context = self._rect_solver._opencl_context self._opencl_last_timing = self._rect_solver._opencl_last_timing return p1, p2 # Note: api.py may call solve_opencl() for single-field OpenCL path; # only solve_opencl_batch2 is implemented. Phase 2 should add or remove that call.