ykcheungab
315 lines
12 KiB
Python
315 lines
12 KiB
Python
import numpy
|
|
from scipy import sparse
|
|
from scipy.sparse import linalg as sparse_linalg
|
|
|
|
from .config import MaskedPoissonConfig
|
|
|
|
|
|
def normalize_mask(mask, shape):
|
|
"""Normalize optional mask to bool array matching grid shape."""
|
|
if mask is None:
|
|
return None
|
|
mask_arr = numpy.asarray(mask, dtype=bool)
|
|
if mask_arr.shape != shape:
|
|
raise ValueError("Mask dimensions should match the Poisson grid")
|
|
return mask_arr
|
|
|
|
|
|
def compute_boundary_mask(active):
|
|
"""Return active pixels with at least one inactive 4-neighbor."""
|
|
up = numpy.zeros_like(active)
|
|
down = numpy.zeros_like(active)
|
|
left = numpy.zeros_like(active)
|
|
right = numpy.zeros_like(active)
|
|
up[1:, :] = active[:-1, :]
|
|
down[:-1, :] = active[1:, :]
|
|
left[:, 1:] = active[:, :-1]
|
|
right[:, :-1] = active[:, 1:]
|
|
fully_surrounded = up & down & left & right
|
|
return active & ~fully_surrounded
|
|
|
|
|
|
def build_masked_factor(lap):
|
|
"""Factorize the masked Laplacian for fast repeated triangular solves.
|
|
|
|
The reduced system is a symmetric M-matrix (graph Laplacian of the active
|
|
pixels with an extra Dirichlet pin), so SuperLU can skip partial pivoting
|
|
(``SymmetricMode`` + ``diag_pivot_thresh=0``), which is ~2x faster than the
|
|
default pivoting factorization. Factoring in float32 halves the memory
|
|
traffic of the triangular solves; relative error against the float64
|
|
reference stays around 1e-7 for these systems.
|
|
|
|
``rhs_dtype`` is returned because ``SuperLU.solve`` refuses to cast a
|
|
float64 right-hand side down to a float32 factor.
|
|
|
|
lap: (N, N) symmetric sparse Laplacian (any scipy dtype)
|
|
returns: (solve_callable, rhs_dtype)
|
|
"""
|
|
lap_csc = lap.tocsc()
|
|
lap32 = lap_csc.astype(numpy.float32, copy=False)
|
|
for diag_pivot_thresh in (0.0, 1.0):
|
|
try:
|
|
lu = sparse_linalg.splu(
|
|
lap32,
|
|
permc_spec="MMD_AT_PLUS_A",
|
|
diag_pivot_thresh=diag_pivot_thresh,
|
|
options=dict(SymmetricMode=True),
|
|
)
|
|
return lu.solve, numpy.float32
|
|
except Exception: # noqa: S112 - try the next pivoting strategy
|
|
# Probe: try the next pivoting strategy; the float64 fallback below is
|
|
# returned if every candidate fails, so this failure is not lost.
|
|
continue
|
|
# Last resort: the original pivoting LU on the float64 matrix.
|
|
return sparse_linalg.factorized(lap_csc), numpy.float64
|
|
|
|
|
|
def sample_boundary_coords(boundary_coords, config: MaskedPoissonConfig):
|
|
"""Subsample boundary pixel coordinates for capacitance correction."""
|
|
if boundary_coords.size == 0:
|
|
return boundary_coords
|
|
|
|
sampled = boundary_coords
|
|
if config.boundary_stride > 1:
|
|
sampled = sampled[:: config.boundary_stride]
|
|
|
|
if config.max_boundary > 0 and len(sampled) > config.max_boundary:
|
|
picks = numpy.linspace(0, len(sampled) - 1, num=config.max_boundary, dtype=numpy.int32)
|
|
sampled = sampled[picks]
|
|
|
|
return numpy.ascontiguousarray(sampled, dtype=numpy.int32)
|
|
|
|
|
|
class MaskedDomain:
|
|
"""Reduced Laplacian system on active mask pixels (2D only)."""
|
|
|
|
__slots__ = (
|
|
"_rhs_buffer",
|
|
"_rhs_pair_buffer",
|
|
"_scatter_buffer",
|
|
"_scatter_pair_buffer",
|
|
"active_coords",
|
|
"active_linear_indices",
|
|
"factor",
|
|
"factor_dtype",
|
|
"index_map",
|
|
"lap",
|
|
"mask",
|
|
"num_active",
|
|
"pinned",
|
|
"shape",
|
|
)
|
|
|
|
def __init__(
|
|
self,
|
|
shape,
|
|
mask,
|
|
index_map,
|
|
factor,
|
|
lap,
|
|
num_active,
|
|
pinned,
|
|
active_coords,
|
|
active_linear_indices,
|
|
factor_dtype=numpy.float32,
|
|
):
|
|
self.shape = shape
|
|
self.mask = mask
|
|
self.index_map = index_map
|
|
self.factor = factor
|
|
self.factor_dtype = numpy.dtype(factor_dtype)
|
|
self.lap = lap
|
|
self.num_active = num_active
|
|
self.pinned = pinned
|
|
self.active_coords = active_coords
|
|
self.active_linear_indices = active_linear_indices
|
|
self._rhs_buffer = None
|
|
self._rhs_pair_buffer = None
|
|
self._scatter_buffer = None
|
|
self._scatter_pair_buffer = None
|
|
|
|
@classmethod
|
|
def build(cls, mask, spacings):
|
|
"""Build masked Laplacian domain from bool mask and grid spacings (dy, dx).
|
|
|
|
mask: (H, W) bool; True = active domain pixels
|
|
spacings: (dy, dx) grid spacing tuple
|
|
returns: MaskedDomain with factorized sparse Laplacian
|
|
"""
|
|
active = numpy.asarray(mask, dtype=bool)
|
|
h, w = active.shape
|
|
active_coords = numpy.argwhere(active)
|
|
num_active = int(active.sum())
|
|
if num_active == 0:
|
|
return cls(
|
|
shape=active.shape,
|
|
mask=active.copy(),
|
|
index_map=numpy.full(active.shape, -1, dtype=numpy.int32),
|
|
factor=None,
|
|
lap=None,
|
|
num_active=0,
|
|
pinned=0,
|
|
active_coords=active_coords,
|
|
active_linear_indices=numpy.flatnonzero(active),
|
|
)
|
|
|
|
index_map = numpy.full(active.shape, -1, dtype=numpy.int32)
|
|
index_map[active] = numpy.arange(num_active, dtype=numpy.int32)
|
|
active_linear_indices = numpy.flatnonzero(active)
|
|
|
|
dy, dx = float(spacings[0]), float(spacings[1])
|
|
coeff_y = 1.0 / (dy * dy)
|
|
coeff_x = 1.0 / (dx * dx)
|
|
|
|
rows = []
|
|
cols = []
|
|
data = []
|
|
pinned = 0
|
|
for row, col in active_coords:
|
|
idx = int(index_map[row, col])
|
|
if idx == pinned:
|
|
rows.append(idx)
|
|
cols.append(idx)
|
|
data.append(1.0)
|
|
continue
|
|
|
|
diag = 0.0
|
|
for drow, dcol, coeff in ((-1, 0, coeff_y), (1, 0, coeff_y), (0, -1, coeff_x), (0, 1, coeff_x)):
|
|
nrow = row + drow
|
|
ncol = col + dcol
|
|
if nrow < 0 or nrow >= h or ncol < 0 or ncol >= w or not active[nrow, ncol]:
|
|
continue
|
|
rows.append(idx)
|
|
cols.append(int(index_map[nrow, ncol]))
|
|
data.append(-coeff)
|
|
diag += coeff
|
|
|
|
if diag == 0.0:
|
|
rows.append(idx)
|
|
cols.append(idx)
|
|
data.append(1.0)
|
|
else:
|
|
rows.append(idx)
|
|
cols.append(idx)
|
|
data.append(diag)
|
|
|
|
lap = sparse.csc_matrix((data, (rows, cols)), shape=(num_active, num_active), dtype=numpy.float64)
|
|
factor, factor_dtype = build_masked_factor(lap)
|
|
return cls(
|
|
shape=active.shape,
|
|
mask=active.copy(),
|
|
index_map=index_map,
|
|
factor=factor,
|
|
lap=lap,
|
|
num_active=num_active,
|
|
pinned=pinned,
|
|
active_coords=active_coords,
|
|
active_linear_indices=active_linear_indices,
|
|
factor_dtype=factor_dtype,
|
|
)
|
|
|
|
def matches(self, mask):
|
|
"""True if this domain was built for the same mask array."""
|
|
mask_arr = numpy.asarray(mask, dtype=bool)
|
|
return self.shape == mask_arr.shape and numpy.array_equal(self.mask, mask_arr)
|
|
|
|
def prepare_rhs(self, f):
|
|
"""Extract gauge-fixed RHS on active pixels.
|
|
|
|
f: (H, W) scalar source field
|
|
returns: (N,) RHS on active pixels; empty if no active pixels
|
|
|
|
The returned dtype matches the factorization (see ``factor_dtype``):
|
|
``SuperLU.solve`` will not downcast a float64 RHS for a float32 factor.
|
|
"""
|
|
if self.num_active == 0:
|
|
return numpy.zeros(0, dtype=self.factor_dtype)
|
|
if (
|
|
self._rhs_buffer is None
|
|
or self._rhs_buffer.shape != (self.num_active,)
|
|
or self._rhs_buffer.dtype != self.factor_dtype
|
|
):
|
|
self._rhs_buffer = numpy.empty(self.num_active, dtype=self.factor_dtype)
|
|
numpy.take(
|
|
numpy.asarray(f, dtype=self.factor_dtype).ravel(),
|
|
self.active_linear_indices,
|
|
out=self._rhs_buffer,
|
|
)
|
|
rhs = self._rhs_buffer
|
|
rhs -= rhs.mean()
|
|
rhs[self.pinned] = 0.0
|
|
return rhs
|
|
|
|
def prepare_rhs_pair(self, f0, f1):
|
|
"""Extract two gauge-fixed RHS vectors into one reusable (N, 2) buffer."""
|
|
if self.num_active == 0:
|
|
return numpy.zeros((0, 2), dtype=self.factor_dtype)
|
|
if (
|
|
self._rhs_pair_buffer is None
|
|
or self._rhs_pair_buffer.shape != (self.num_active, 2)
|
|
or self._rhs_pair_buffer.dtype != self.factor_dtype
|
|
):
|
|
self._rhs_pair_buffer = numpy.empty((self.num_active, 2), dtype=self.factor_dtype)
|
|
flat0 = numpy.asarray(f0, dtype=self.factor_dtype).ravel()
|
|
flat1 = numpy.asarray(f1, dtype=self.factor_dtype).ravel()
|
|
numpy.take(flat0, self.active_linear_indices, out=self._rhs_pair_buffer[:, 0])
|
|
numpy.take(flat1, self.active_linear_indices, out=self._rhs_pair_buffer[:, 1])
|
|
self._rhs_pair_buffer -= self._rhs_pair_buffer.mean(axis=0, keepdims=True)
|
|
self._rhs_pair_buffer[self.pinned, :] = 0.0
|
|
return self._rhs_pair_buffer
|
|
|
|
def scatter(self, sol_active, template):
|
|
"""Map reduced solution back to full grid.
|
|
|
|
sol_active: (N,) potential on active pixels
|
|
template: (H, W) array used for output shape/dtype
|
|
returns: (H, W) potential; zero outside mask
|
|
"""
|
|
if self._scatter_buffer is None or self._scatter_buffer.shape != template.shape:
|
|
self._scatter_buffer = numpy.zeros(template.shape, dtype=numpy.float32)
|
|
else:
|
|
self._scatter_buffer.fill(0)
|
|
if sol_active.size:
|
|
self._scatter_buffer.ravel()[self.active_linear_indices] = numpy.asarray(sol_active, dtype=numpy.float32)
|
|
return self._scatter_buffer.copy()
|
|
|
|
def scatter_pair(self, sol_active_pair, template):
|
|
"""Map two reduced solutions back to a pair of full-grid outputs."""
|
|
if self._scatter_pair_buffer is None or self._scatter_pair_buffer.shape != (2,) + template.shape:
|
|
self._scatter_pair_buffer = numpy.zeros((2,) + template.shape, dtype=numpy.float32)
|
|
else:
|
|
self._scatter_pair_buffer.fill(0)
|
|
if sol_active_pair.size:
|
|
values = numpy.asarray(sol_active_pair, dtype=numpy.float32)
|
|
self._scatter_pair_buffer[0].ravel()[self.active_linear_indices] = values[:, 0]
|
|
self._scatter_pair_buffer[1].ravel()[self.active_linear_indices] = values[:, 1]
|
|
return self._scatter_pair_buffer[0].copy(), self._scatter_pair_buffer[1].copy()
|
|
|
|
def residual(self, field, rhs_active):
|
|
"""Compute masked Laplacian residual L u - rhs on active pixels.
|
|
|
|
field: (H, W) full-grid field
|
|
rhs_active: (N,) RHS on active pixels (or None for homogeneous)
|
|
returns: (N,) residual vector
|
|
"""
|
|
vec = numpy.asarray(field, dtype=numpy.float64)[self.mask]
|
|
residual = self.lap.dot(vec)
|
|
if rhs_active is not None:
|
|
residual = residual - rhs_active
|
|
return residual
|
|
|
|
def to_cache_dict(self, capfft=None):
|
|
"""Legacy dict view for tests and facade compatibility."""
|
|
return {
|
|
"shape": self.shape,
|
|
"mask": self.mask,
|
|
"index_map": self.index_map,
|
|
"factor": self.factor,
|
|
"factor_dtype": self.factor_dtype,
|
|
"lap": self.lap,
|
|
"num_active": self.num_active,
|
|
"pinned": self.pinned,
|
|
"active_coords": self.active_coords,
|
|
"capfft": capfft,
|
|
}
|