Speed up rasterization ~40x and large solves ~2x

- Hybrid rasterizer: PIL scanline fill for the bulk, with cells in a
  ~2 px band around each ring edge re-tested exactly against the
  polygon - cell-for-cell identical to the old center-in-polygon pass
  (equivalence test added) but O(vertices + cells) instead of
  O(vertices x cells). Measured 4.5 s -> 0.11 s at 1.45M cells with
  8.8k polygon vertices.
- AMG-preconditioned CG (pyamg, new requirement) above 500k unknowns:
  measured 7.0 s vs 15.3 s spsolve at 1.4M unknowns at a fraction of
  the memory, R identical to 1e-6; the old Jacobi-CG (kept as fallback
  when pyamg is missing) needed tens of minutes there. spsolve stays
  the default below 500k where it is exact and fastest.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
janik
2026-07-15 15:24:21 +07:00
parent 6e989fc5f8
commit b62e45a9b4
9 changed files with 124 additions and 23 deletions
+30
View File
@@ -293,10 +293,40 @@ def solve_system(A: sparse.csr_matrix, b: np.ndarray) -> tuple[np.ndarray, Solve
if n <= config.SPSOLVE_MAX_UNKNOWNS:
x = sla.spsolve(A.tocsc(), b)
return x, SolveInfo(method="spsolve", n_unknowns=n)
try:
return _solve_amg(A, b)
except ImportError:
print("note: pyamg not installed - falling back to Jacobi-CG "
"(much slower on large grids)")
return _solve_cg_jacobi(A, b)
def _solve_amg(A: sparse.csr_matrix, b: np.ndarray) -> tuple[np.ndarray, SolveInfo]:
"""CG preconditioned with smoothed-aggregation AMG: near-linear
scaling on these 2D Laplacians and a fraction of spsolve's memory."""
import pyamg
n = A.shape[0]
ml = pyamg.smoothed_aggregation_solver(A.tocsr(), max_coarse=500)
residuals: list[float] = []
x = ml.solve(b, tol=config.AMG_TOL, maxiter=300, accel="cg",
residuals=residuals)
res = float(np.linalg.norm(b - A @ x) / max(np.linalg.norm(b), 1e-300))
if not np.isfinite(res) or res > 1e-6:
raise SolverError(
f"AMG-CG did not converge (residual {res:.2e}). Try a "
f"different grid size, or force the direct solver by raising "
f"SPSOLVE_MAX_UNKNOWNS in config.py."
)
return x, SolveInfo(method="amg+cg", n_unknowns=n,
iterations=max(len(residuals) - 1, 0), residual=res)
def _solve_cg_jacobi(A: sparse.csr_matrix, b: np.ndarray) -> tuple[np.ndarray, SolveInfo]:
# The matrix is SPD, so CG is guaranteed to converge. Jacobi is the
# only preconditioner in scipy that keeps the preconditioned operator
# SPD without a factorization that can break down at this scale.
n = A.shape[0]
d = A.diagonal()
M = sla.LinearOperator((n, n), lambda v: v / d)
iters = 0