integrated finger process pipeline on remapped 2d image

2026-06-14 17:48:30 +08:00
parent 667bdf1961
commit 30bebcd439
36 changed files with 3593 additions and 2043 deletions
@@ -0,0 +1,113 @@
+import numpy
+from scipy import linalg as dense_linalg
+
+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 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_solver.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_solver.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_solver.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)