Skip to content

Commit 086c727

Browse files
author
James Nightingale
authored
Merge pull request #652 from PyAutoLabs/feature/matrix-free-interferometer-phase4
feat: matrix-free interferometer positivity via primal-dual interior point with active-set polish (phase 4 of #638)
2 parents 625e119 + 850ea01 commit 086c727

2 files changed

Lines changed: 557 additions & 11 deletions

File tree

‎autoarray/inversion/inversion/interferometer/inversion_interferometer_matrix_free_util.py‎

Lines changed: 342 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -29,10 +29,30 @@
2929
`W~` by its peak, which keeps it sparse; it is factorised once per `λ` by a sparse LU.
3030
3131
Phase 1 landed the operator data term and the solver as library utilities, phase 2 the natively sparse
32-
Matérn precision (`aa.reg.MaternSparse`), and phase 3 the evidence log-determinants (below), each with
33-
parity tests against the dense sparse-path quantities. None of it is wired into fits and there is no
34-
`Settings` flag. Later phases (interior-point positivity, profiling and workspace examples) follow as
35-
their own tasks.
32+
Matérn precision (`aa.reg.MaternSparse`), phase 3 the evidence log-determinants (below) and phase 4 the
33+
positivity constraint (below), each with parity tests against the dense sparse-path quantities. None of
34+
it is wired into fits and there is no `Settings` flag.
35+
36+
Positivity (phase 4)
37+
--------------------
38+
`reconstruction_positive_via_interior_point_from` solves `(F + λH) s = D` subject to `s >= 0` by a
39+
Mehrotra predictor-corrector primal-dual interior point method whose Newton systems are solved by CG
40+
preconditioned by the Powell matrix plus the barrier diagonal, followed by an exact solve on the
41+
detected support with primal-dual active-set corrections (Hintermüller, Ito & Kunisch 2002), so the
42+
held pixels are exact zeros. On the test case (a `10 x 10` `MaternSparse` mesh on the data-dominated
43+
dataset, 33-72% of pixels at zero) it reproduces the support of `scipy.optimize.nnls` on the whitened
44+
dense system exactly and its values to ~1e-10 of the max in 8-10 interior-point steps and one polish
45+
solve.
46+
47+
Remaining work (phase 5)
48+
------------------------
49+
- wiring as an inversion class or a `Settings` option, so fits can select the matrix-free path;
50+
- a profiling note, dense vs matrix-free at `50^2`, `64^2`, `128^2` and `192^2` mesh pixels, ingested
51+
by PyAutoPulse;
52+
- a workspace example;
53+
- reconstruction uncertainty maps, the diagonal of `(F + λH)⁻¹`, by a stochastic diagonal estimator
54+
(there is no dense inverse to read it from);
55+
- linear light profiles and multiple mappers (every function here assumes a single mapper's `M`).
3656
3757
The log evidence (phase 3)
3858
--------------------------
@@ -381,8 +401,8 @@ def reconstruction_via_pcg_from(
381401
preconditioner (`powell_preconditioner_from`) applied through its sparse LU solve.
382402
383403
This is the matrix-free counterpart of the positive-negative (`use_positive_only_solver=False`)
384-
solve of an inversion's `curvature_reg_matrix`. No positivity is imposed (an interior-point
385-
positive solve is a later phase).
404+
solve of an inversion's `curvature_reg_matrix`. No positivity is imposed (see
405+
`reconstruction_positive_via_interior_point_from` for the positive solve).
386406
387407
Parameters
388408
----------
@@ -731,3 +751,319 @@ def log_evidence_matrix_free_from(
731751
+ float(noise_normalization)
732752
)
733753
)
754+
755+
756+
class PositiveReconstructionResult:
757+
def __init__(
758+
self,
759+
reconstruction: np.ndarray,
760+
support: np.ndarray,
761+
interior_point_steps: int,
762+
polish_iterations: int,
763+
gap: float,
764+
residual: float,
765+
converged: bool,
766+
):
767+
"""
768+
The output of `reconstruction_positive_via_interior_point_from`.
769+
770+
Parameters
771+
----------
772+
reconstruction
773+
The non-negative reconstruction `s`, with exact zeros on the held (active) set.
774+
support
775+
Boolean mask of the free (passive) pixels, those solved for; `reconstruction` is zero
776+
wherever it is `False`.
777+
interior_point_steps
778+
The number of Mehrotra predictor-corrector steps taken.
779+
polish_iterations
780+
The number of exact free-set solves in the active-set polish (one if the support detected
781+
from the interior point was already optimal).
782+
gap
783+
The normalised duality gap `n μ / |xᵀD|` of the final interior-point iterate.
784+
residual
785+
The normalised dual residual `|A x - D - z| / |D|` of the final interior-point iterate.
786+
converged
787+
`True` if the active-set polish ended on a support satisfying the KKT conditions (no
788+
negative free pixel and no held pixel with a negative gradient), `False` if it reached
789+
its iteration cap and the result was clipped at zero.
790+
"""
791+
self.reconstruction = reconstruction
792+
self.support = support
793+
self.interior_point_steps = interior_point_steps
794+
self.polish_iterations = polish_iterations
795+
self.gap = gap
796+
self.residual = residual
797+
self.converged = converged
798+
799+
800+
def _step_to_boundary(v: np.ndarray, dv: np.ndarray) -> float:
801+
"""
802+
The largest step `α <= 1` keeping `v + α dv >= 0` (for `v > 0`).
803+
"""
804+
decreasing = dv < 0.0
805+
806+
if not np.any(decreasing):
807+
return 1.0
808+
809+
return float(min(1.0, np.min(-v[decreasing] / dv[decreasing])))
810+
811+
812+
def reconstruction_positive_via_interior_point_from(
813+
operator: MatrixFreeDataOperator,
814+
regularization_matrix_sparse: scipy.sparse.spmatrix,
815+
coefficient: float,
816+
max_steps: int = 60,
817+
tol: float = 1.0e-9,
818+
cg_rtol_min: float = 1.0e-8,
819+
cg_rtol_max: float = 1.0e-2,
820+
polish_rtol: float = 1.0e-10,
821+
max_polish: int = 10,
822+
maxiter: int = 5000,
823+
) -> PositiveReconstructionResult:
824+
"""
825+
Returns the non-negative reconstruction solving
826+
827+
min_s ½ sᵀ(F + λH)s - sᵀD subject to s >= 0
828+
829+
(equivalently `(F + λH) s = D` on the free pixels, zero elsewhere) without forming or factorising
830+
the dense `(S, S)` matrix `A = F + λH`: every operation is a matrix-free product with `A` or a
831+
sparse factorisation of the Powell preconditioner `P = κ₀ MᵀM + λH` plus a diagonal.
832+
833+
This is the matrix-free counterpart of the positive-only (`use_positive_only_solver=True`) solve
834+
of an inversion's `curvature_reg_matrix`. The method follows @HRSAstro's reference
835+
(PyAutoArray #638, phase 4) in two stages.
836+
837+
**Primal-dual interior point (Mehrotra predictor-corrector).** The KKT conditions are
838+
`A x - D = z`, `x >= 0`, `z >= 0`, `x z = 0`. Starting from one Jacobi step,
839+
`x₀ = max(D / diag, 0)` shifted by `10⁻² max(x₀)` into the interior (with `diag` the diagonal of
840+
`P`, the Powell approximation of the diagonal of `A`), and a dual `z₀` from the positive part of
841+
the gradient, each step eliminates `dz` from the linearised KKT system, leaving
842+
843+
(A + Z X⁻¹) dx = -r_d - z + c, r_d = A x - D - z,
844+
845+
with `c = 0` for the affine (predictor) direction and `c = (σμ - dx_aff dz_aff) / x` for the
846+
combined direction, centring parameter `σ = (μ_aff / μ)³` and `μ = xᵀz / n`. Each of the two
847+
solves is CG preconditioned by the barrier-diagonal-augmented Powell matrix `P + Z X⁻¹`, refactorised
848+
by a sparse LU every step (the barrier diagonal changes); the corrector is warm-started at the
849+
predictor. The CG tolerance is tied to the progress, `0.1 min(gap, residual + gap)`, clipped to
850+
`[cg_rtol_min, cg_rtol_max]`, so early steps are cheap. Steps are damped to 0.99 of the distance to
851+
the boundary. Iteration stops when the normalised gap `n μ / |xᵀD|` and dual residual
852+
`|r_d| / |D|` are both below `tol`, or after `max_steps`.
853+
854+
**Exact active-set polish.** The interior point never reaches exact zeros, so its support is read
855+
off as the free set `{i : z_i <= diag_i x_i}` and the free pixels are solved exactly,
856+
`A_ff s_f = D_f`, by CG preconditioned by the sparse LU of the sub-matrix `P_ff` and warm-started
857+
at the interior-point iterate; the held pixels are exactly zero. The KKT conditions are then
858+
checked through the multiplier `g = A s - D` and the support updated by the primal-dual active-set
859+
rule (Hintermüller, Ito & Kunisch 2002): a pixel is held when `g_i > diag_i s_i`, which releases
860+
held pixels with a negative gradient and holds free pixels that came out negative. This repeats
861+
until the support is unchanged (normally after one or two solves), or `max_polish` solves, after
862+
which the solution is clipped at zero and `converged` is `False`.
863+
864+
The contributor found projected Newton and a pure active-set method stall at weak priors with
865+
~30% of pixels at zero; the interior point is insensitive to the size of the active set and the
866+
polish starts from an essentially correct support, so the result matches `scipy.optimize.nnls`
867+
on the whitened dense system.
868+
869+
Parameters
870+
----------
871+
operator
872+
The data operator providing `F` and `D`.
873+
regularization_matrix_sparse
874+
The regularization matrix `H` as a `scipy.sparse` matrix.
875+
coefficient
876+
The scalar `λ` multiplying `H`; `1.0` for an autoarray regularization matrix, which already
877+
carries its coefficient.
878+
max_steps
879+
The maximum number of interior-point steps.
880+
tol
881+
The normalised gap and dual residual at which the interior point stops.
882+
cg_rtol_min, cg_rtol_max
883+
The clip range of the interior-point CG tolerance.
884+
polish_rtol
885+
The CG tolerance of the exact free-set solves, relative to `|D_f|`.
886+
max_polish
887+
The maximum number of free-set solves in the active-set polish.
888+
maxiter
889+
The maximum number of iterations of every CG solve.
890+
891+
Returns
892+
-------
893+
A `PositiveReconstructionResult` holding the reconstruction (exact zeros on the held set), the
894+
support mask, the interior-point step and polish iteration counts and the final gap and residual.
895+
"""
896+
regularization_matrix_sparse = scipy.sparse.csr_matrix(regularization_matrix_sparse)
897+
898+
n = operator.params
899+
data_vector = operator.data_vector
900+
data_vector_norm = float(np.linalg.norm(data_vector))
901+
902+
def curvature_reg_matvec(v):
903+
return operator.curvature_matvec(v) + coefficient * (
904+
regularization_matrix_sparse @ v
905+
)
906+
907+
preconditioner, _ = powell_preconditioner_from(
908+
operator=operator,
909+
regularization_matrix_sparse=regularization_matrix_sparse,
910+
coefficient=coefficient,
911+
)
912+
preconditioner = preconditioner.tocsc()
913+
diag = preconditioner.diagonal()
914+
915+
def splu_from(matrix):
916+
return scipy.sparse.linalg.splu(
917+
matrix.tocsc(),
918+
permc_spec="MMD_AT_PLUS_A",
919+
diag_pivot_thresh=0.0,
920+
options=dict(SymmetricMode=True),
921+
)
922+
923+
# Starting point: one Jacobi step, shifted into the interior, and a positive dual.
924+
x = np.clip(data_vector / diag, 0.0, None)
925+
x = x + 1.0e-2 * max(float(np.max(x)), np.finfo(float).tiny)
926+
927+
gradient = curvature_reg_matvec(x) - data_vector
928+
z = np.maximum(gradient, 0.0) + 1.0e-2 * max(
929+
float(np.max(np.abs(gradient))), np.finfo(float).tiny
930+
)
931+
932+
interior_point_steps = 0
933+
gap = residual = np.inf
934+
935+
for _ in range(max_steps):
936+
dual_residual = gradient - z
937+
mu = float(x @ z) / n
938+
939+
gap = n * mu / abs(float(x @ data_vector))
940+
residual = float(np.linalg.norm(dual_residual)) / data_vector_norm
941+
942+
if gap < tol and residual < tol:
943+
break
944+
945+
interior_point_steps += 1
946+
947+
barrier_diagonal = z / x
948+
949+
factor = splu_from(preconditioner + scipy.sparse.diags(barrier_diagonal))
950+
951+
newton_operator = scipy.sparse.linalg.LinearOperator(
952+
(n, n),
953+
matvec=lambda v: curvature_reg_matvec(v) + barrier_diagonal * v,
954+
dtype=np.float64,
955+
)
956+
preconditioner_operator = scipy.sparse.linalg.LinearOperator(
957+
(n, n), matvec=factor.solve, dtype=np.float64
958+
)
959+
960+
rtol = float(np.clip(0.1 * min(gap, residual + gap), cg_rtol_min, cg_rtol_max))
961+
962+
# Predictor (affine scaling) direction.
963+
dx, _ = scipy.sparse.linalg.cg(
964+
newton_operator,
965+
-dual_residual - z,
966+
rtol=rtol,
967+
maxiter=maxiter,
968+
M=preconditioner_operator,
969+
)
970+
dz = -z - barrier_diagonal * dx
971+
972+
mu_affine = (
973+
float(
974+
(x + _step_to_boundary(x, dx) * dx)
975+
@ (z + _step_to_boundary(z, dz) * dz)
976+
)
977+
/ n
978+
)
979+
sigma = (mu_affine / mu) ** 3
980+
981+
# Combined centring + corrector direction, warm-started at the predictor.
982+
correction = (sigma * mu - dx * dz) / x
983+
984+
dx, _ = scipy.sparse.linalg.cg(
985+
newton_operator,
986+
-dual_residual - z + correction,
987+
x0=dx,
988+
rtol=rtol,
989+
maxiter=maxiter,
990+
M=preconditioner_operator,
991+
)
992+
dz = -z + correction - barrier_diagonal * dx
993+
994+
x = x + 0.99 * _step_to_boundary(x, dx) * dx
995+
z = z + 0.99 * _step_to_boundary(z, dz) * dz
996+
997+
gradient = curvature_reg_matvec(x) - data_vector
998+
999+
# Exact active-set polish on the support detected from the interior point.
1000+
free = z <= diag * x
1001+
warm_start = x.copy()
1002+
1003+
reconstruction = np.zeros(n)
1004+
polish_iterations = 0
1005+
converged = False
1006+
1007+
for _ in range(max_polish):
1008+
polish_iterations += 1
1009+
1010+
free_indexes = np.flatnonzero(free)
1011+
reconstruction = np.zeros(n)
1012+
1013+
if free_indexes.size > 0:
1014+
1015+
def free_matvec(v, free_indexes=free_indexes):
1016+
return curvature_reg_matvec(
1017+
np.bincount(free_indexes, weights=v, minlength=n)
1018+
)[free_indexes]
1019+
1020+
sub_factor = splu_from(preconditioner[free_indexes][:, free_indexes])
1021+
1022+
free_operator = scipy.sparse.linalg.LinearOperator(
1023+
(free_indexes.size, free_indexes.size),
1024+
matvec=free_matvec,
1025+
dtype=np.float64,
1026+
)
1027+
sub_preconditioner_operator = scipy.sparse.linalg.LinearOperator(
1028+
(free_indexes.size, free_indexes.size),
1029+
matvec=sub_factor.solve,
1030+
dtype=np.float64,
1031+
)
1032+
1033+
reconstruction_free, _ = scipy.sparse.linalg.cg(
1034+
free_operator,
1035+
data_vector[free_indexes],
1036+
x0=warm_start[free_indexes],
1037+
rtol=polish_rtol,
1038+
maxiter=maxiter,
1039+
M=sub_preconditioner_operator,
1040+
)
1041+
1042+
reconstruction[free_indexes] = reconstruction_free
1043+
1044+
# Primal-dual active-set update: the multiplier of the held pixels is the gradient.
1045+
multiplier = curvature_reg_matvec(reconstruction) - data_vector
1046+
multiplier[free] = 0.0
1047+
1048+
free_new = ~(multiplier > diag * reconstruction)
1049+
1050+
if np.array_equal(free_new, free):
1051+
converged = bool(np.all(reconstruction >= 0.0))
1052+
break
1053+
1054+
free = free_new
1055+
warm_start = np.clip(reconstruction, 0.0, None)
1056+
1057+
if not converged:
1058+
reconstruction = np.clip(reconstruction, 0.0, None)
1059+
free = reconstruction > 0.0
1060+
1061+
return PositiveReconstructionResult(
1062+
reconstruction=reconstruction,
1063+
support=free,
1064+
interior_point_steps=interior_point_steps,
1065+
polish_iterations=polish_iterations,
1066+
gap=float(gap),
1067+
residual=float(residual),
1068+
converged=converged,
1069+
)

0 commit comments

Comments
 (0)