ykcheungab
134 lines
5.5 KiB
Python
134 lines
5.5 KiB
Python
import numpy
|
|
from scipy import linalg as dense_linalg
|
|
|
|
from ...backend_availability import CUPY_AVAILABLE, cp
|
|
from ..config import MaskedPoissonConfig
|
|
from ..domain import MaskedDomain, compute_boundary_mask, sample_boundary_coords
|
|
|
|
|
|
class CapacitanceFFTSolver:
|
|
"""Masked Poisson via rectangular FFT bulk solve + capacitance boundary correction."""
|
|
|
|
def __init__(self, config: MaskedPoissonConfig, bulk_solver):
|
|
self.config = config
|
|
self.bulk_solver = bulk_solver
|
|
self._capfft_cache = None
|
|
self._capfft_domain = None
|
|
|
|
def _use_cupy_backend(self):
|
|
return (
|
|
getattr(self.bulk_solver, "poisson_backend", "cpu") == "cupy"
|
|
and CUPY_AVAILABLE
|
|
)
|
|
|
|
def _bulk_solve(self, rhs):
|
|
"""Solve one full-grid Poisson problem, using CuPy when configured."""
|
|
if not self._use_cupy_backend():
|
|
return self.bulk_solver.solve(rhs)
|
|
|
|
if not getattr(self.bulk_solver, "ready", False):
|
|
self.bulk_solver.prepare(False)
|
|
|
|
rhs_gpu = cp.asarray(rhs, dtype=cp.float32)
|
|
field_gpu = self.bulk_solver._solve_gpu_internal(rhs_gpu)
|
|
cp.cuda.get_current_stream().synchronize()
|
|
return cp.asnumpy(field_gpu)
|
|
|
|
def prepare(self, domain: MaskedDomain):
|
|
"""Precompute capacitance matrix for boundary correction.
|
|
|
|
domain: MaskedDomain with built Laplacian
|
|
returns: capacitance cache dict (also stored internally)
|
|
"""
|
|
self._capfft_domain = domain
|
|
self._capfft_cache = self._build_capfft_cache(domain)
|
|
return self._capfft_cache
|
|
|
|
def _build_capfft_cache(self, domain: MaskedDomain):
|
|
active = domain.mask
|
|
boundary_coords = numpy.argwhere(compute_boundary_mask(active))
|
|
boundary_coords = sample_boundary_coords(boundary_coords, self.config)
|
|
if boundary_coords.size == 0:
|
|
pinned_rc = numpy.array(domain.active_coords[domain.pinned], dtype=numpy.int32)
|
|
boundary_coords = pinned_rc.reshape(1, 2)
|
|
|
|
boundary_active_indices = domain.index_map[boundary_coords[:, 0], boundary_coords[:, 1]]
|
|
boundary_count = int(len(boundary_active_indices))
|
|
|
|
cap_matrix = numpy.zeros((boundary_count, boundary_count), dtype=numpy.float64)
|
|
zero_rhs = numpy.zeros(domain.num_active, dtype=numpy.float64)
|
|
for col, (row, col_idx) in enumerate(boundary_coords):
|
|
charge = numpy.zeros(active.shape, dtype=numpy.float32)
|
|
charge[row, col_idx] = 1.0
|
|
response = numpy.asarray(self._bulk_solve(charge), dtype=numpy.float64)
|
|
residual = domain.residual(response, zero_rhs)
|
|
cap_matrix[:, col] = residual[boundary_active_indices]
|
|
|
|
reg_scale = self.config.reg
|
|
if reg_scale <= 0:
|
|
reg_scale = 1e-8
|
|
reg = reg_scale * max(1.0, float(numpy.linalg.norm(cap_matrix, ord=numpy.inf)))
|
|
cap_matrix_reg = cap_matrix + reg * numpy.eye(boundary_count, dtype=numpy.float64)
|
|
lu, piv = dense_linalg.lu_factor(cap_matrix_reg)
|
|
|
|
return {
|
|
"boundary_coords": boundary_coords,
|
|
"boundary_active_indices": boundary_active_indices,
|
|
"boundary_count": boundary_count,
|
|
"cap_matrix": cap_matrix,
|
|
"cap_matrix_reg": cap_matrix_reg,
|
|
"cap_lu": lu,
|
|
"cap_piv": piv,
|
|
"reg": reg,
|
|
}
|
|
|
|
def get_cache(self, domain: MaskedDomain):
|
|
"""Return cached capacitance data, building if needed."""
|
|
if self._capfft_cache is None or self._capfft_domain is not domain:
|
|
self.prepare(domain)
|
|
return self._capfft_cache
|
|
|
|
def solve(self, f, domain: MaskedDomain):
|
|
"""Solve masked Poisson with capacitance-corrected rectangular FFT.
|
|
|
|
f: (H, W) scalar source field
|
|
domain: MaskedDomain with prebuilt Laplacian
|
|
returns: (H, W) scalar potential; zero outside mask
|
|
"""
|
|
rhs_active = domain.prepare_rhs(f)
|
|
if rhs_active.size == 0:
|
|
return numpy.zeros_like(f, dtype=numpy.float32)
|
|
|
|
rhs_full = numpy.zeros_like(f, dtype=numpy.float32)
|
|
rhs_full[domain.mask] = rhs_active.astype(numpy.float32)
|
|
|
|
field = numpy.asarray(self._bulk_solve(rhs_full), dtype=numpy.float64)
|
|
cap_cache = self.get_cache(domain)
|
|
if not cap_cache or cap_cache["boundary_count"] == 0:
|
|
field_active = field[domain.mask]
|
|
if field_active.size:
|
|
field_active = field_active.copy()
|
|
field_active -= field_active[domain.pinned]
|
|
return domain.scatter(field_active, f)
|
|
|
|
boundary_rows = cap_cache["boundary_active_indices"]
|
|
rhs_scale = max(1.0, float(numpy.linalg.norm(rhs_active)))
|
|
for _ in range(self.config.max_iters):
|
|
residual = domain.residual(field, rhs_active)
|
|
boundary_residual = residual[boundary_rows]
|
|
if float(numpy.linalg.norm(boundary_residual)) <= self.config.tol * rhs_scale:
|
|
break
|
|
|
|
alpha = dense_linalg.lu_solve((cap_cache["cap_lu"], cap_cache["cap_piv"]), boundary_residual)
|
|
charges = numpy.zeros_like(f, dtype=numpy.float32)
|
|
coords = cap_cache["boundary_coords"]
|
|
charges[coords[:, 0], coords[:, 1]] = alpha.astype(numpy.float32)
|
|
correction = numpy.asarray(self._bulk_solve(charges), dtype=numpy.float64)
|
|
field -= correction
|
|
|
|
field_active = field[domain.mask]
|
|
if field_active.size:
|
|
field_active = field_active.copy()
|
|
field_active -= field_active[domain.pinned]
|
|
return domain.scatter(field_active, f)
|