|
29 | 29 | `W~` by its peak, which keeps it sparse; it is factorised once per `λ` by a sparse LU. |
30 | 30 |
|
31 | 31 | 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`). |
36 | 56 |
|
37 | 57 | The log evidence (phase 3) |
38 | 58 | -------------------------- |
@@ -381,8 +401,8 @@ def reconstruction_via_pcg_from( |
381 | 401 | preconditioner (`powell_preconditioner_from`) applied through its sparse LU solve. |
382 | 402 |
|
383 | 403 | 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). |
386 | 406 |
|
387 | 407 | Parameters |
388 | 408 | ---------- |
@@ -731,3 +751,319 @@ def log_evidence_matrix_free_from( |
731 | 751 | + float(noise_normalization) |
732 | 752 | ) |
733 | 753 | ) |
| 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