tuned pipeline to 60 fps
This commit is contained in:
@@ -1,6 +1,7 @@
|
||||
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
|
||||
|
||||
@@ -14,6 +15,25 @@ class CapacitanceFFTSolver:
|
||||
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.
|
||||
|
||||
@@ -40,7 +60,7 @@ class CapacitanceFFTSolver:
|
||||
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)
|
||||
response = numpy.asarray(self._bulk_solve(charge), dtype=numpy.float64)
|
||||
residual = domain.residual(response, zero_rhs)
|
||||
cap_matrix[:, col] = residual[boundary_active_indices]
|
||||
|
||||
@@ -82,7 +102,7 @@ class CapacitanceFFTSolver:
|
||||
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)
|
||||
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]
|
||||
@@ -103,7 +123,7 @@ class CapacitanceFFTSolver:
|
||||
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)
|
||||
correction = numpy.asarray(self._bulk_solve(charges), dtype=numpy.float64)
|
||||
field -= correction
|
||||
|
||||
field_active = field[domain.mask]
|
||||
|
||||
@@ -19,3 +19,18 @@ class MaskedSparseSolver:
|
||||
|
||||
sol_active = domain.factor(rhs) # (N,) reduced-system potential
|
||||
return domain.scatter(sol_active, f)
|
||||
|
||||
def solve_pair(self, f0, f1, domain: MaskedDomain):
|
||||
"""Solve two masked Poisson RHS with one reused sparse factorization call."""
|
||||
rhs_pair = domain.prepare_rhs_pair(f0, f1)
|
||||
if rhs_pair.shape[0] == 0:
|
||||
empty = numpy.zeros_like(f0, dtype=numpy.float32)
|
||||
return empty, empty.copy()
|
||||
|
||||
try:
|
||||
sol_pair = domain.factor(rhs_pair) # (N, 2) reduced-system potentials
|
||||
except Exception:
|
||||
sol_pair = numpy.column_stack(
|
||||
[domain.factor(rhs_pair[:, 0]), domain.factor(rhs_pair[:, 1])]
|
||||
)
|
||||
return domain.scatter_pair(sol_pair, f0)
|
||||
|
||||
Reference in New Issue
Block a user