ykcheungab
40 lines
1.5 KiB
Python
40 lines
1.5 KiB
Python
import numpy
|
|
|
|
from ..domain import MaskedDomain
|
|
|
|
|
|
class MaskedSparseSolver:
|
|
"""Reference masked Poisson solver via factorized sparse Laplacian."""
|
|
|
|
def solve(self, f, domain: MaskedDomain):
|
|
"""Solve masked Poisson on active domain pixels.
|
|
|
|
f: (H, W) scalar source field
|
|
domain: MaskedDomain with prebuilt Laplacian
|
|
returns: (H, W) scalar potential; zero outside mask
|
|
"""
|
|
rhs = domain.prepare_rhs(f)
|
|
if rhs.size == 0:
|
|
return numpy.zeros_like(f, dtype=numpy.float32)
|
|
|
|
# MaskedDomain builds a positive diagonal / negative off-diagonal stencil
|
|
# for (-Laplace). Negate the direct solve result so the masked sparse path
|
|
# matches the FFT solver convention used elsewhere: Laplace(u) = rhs.
|
|
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)
|