ykcheungab
358 lines
14 KiB
Python
358 lines
14 KiB
Python
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.
|