Skip to content

Commit 02aa2e4

Browse files
Jammy2211claude
andcommitted
perf: reuse one Cholesky factor of F+H for the solve and the log det
On the NumPy backend with use_positive_only_solver=False, a likelihood call factorized the same SPD matrix F + H twice: np.linalg.solve (LU) for the reconstruction and np.linalg.cholesky for log_det_curvature_reg_matrix_term. A new cached_property `_curvature_reg_matrix_cho_factor` computes one scipy.linalg.cho_factor (a copy; curvature_reg_matrix is not overwritten). `reconstruction` solves with cho_solve, and the log det reads 2 * sum(log(diag(L))) off the same factor when every linear object is regularized. The factor is None (unchanged routes) on JAX, with the positive-only solver, under log_det_method="slogdet", on LinAlgError, or when its diagonal is not finite. An indefinite-but-nonsingular F + H is therefore still solved by the LU and still makes the log det raise (or return 0.0 in test mode), and a singular one still makes the solve raise (or return its test-mode dummy). Tests pin all of these against the two-factorization values. Closes #633 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
1 parent 17caf22 commit 02aa2e4

2 files changed

Lines changed: 362 additions & 1 deletion

File tree

‎autoarray/inversion/inversion/abstract.py‎

Lines changed: 82 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -761,12 +761,80 @@ def reconstruction(self) -> np.ndarray:
761761

762762
return reconstruction
763763

764+
# One Cholesky factor of `F + H`, shared with `log_det_curvature_reg_matrix_term`. `None`
765+
# when the gate in `_curvature_reg_matrix_cho_factor` is closed or the factorization failed,
766+
# in which case the unchanged LU solve below runs (with its test-mode guard).
767+
cho_factor = self._curvature_reg_matrix_cho_factor
768+
769+
if cho_factor is not None:
770+
import scipy.linalg
771+
772+
return scipy.linalg.cho_solve(
773+
cho_factor, np.asarray(self.data_vector), check_finite=False
774+
)
775+
764776
return inversion_util.reconstruction_positive_negative_from(
765777
data_vector=self.data_vector,
766778
curvature_reg_matrix=self.curvature_reg_matrix,
767779
xp=self._xp,
768780
)
769781

782+
@cached_property
783+
def _curvature_reg_matrix_cho_factor(self):
784+
"""
785+
One Cholesky factor of ``F + reg_coeff*H`` for the positive-negative solve, shared by
786+
:attr:`reconstruction` and :attr:`log_det_curvature_reg_matrix_term`.
787+
788+
Without it, a NumPy positive-negative likelihood call factorizes the same symmetric
789+
positive-definite matrix twice: ``np.linalg.solve`` (an LU) for the reconstruction and
790+
``np.linalg.cholesky`` for the log determinant. On the Aris SPT2319-55 interferometer setup
791+
those two were 31% of a call at 1,405 source pixels and 72% at 2,947 (PyAutoArray#633,
792+
campaign PyAutoPulse#62). One ``cho_factor`` serves both: ``cho_solve`` gives the
793+
reconstruction and ``2 * sum(log(diag(L)))`` the log determinant.
794+
795+
Returns the ``(c, lower)`` tuple of :func:`scipy.linalg.cho_factor`, or ``None`` -- meaning
796+
"take the unchanged routes" -- when any of these holds:
797+
798+
- the JAX backend (its solve and log determinant are traced and stay as they are);
799+
- ``Settings.use_positive_only_solver`` (that solve builds and reuses its own factor);
800+
- ``Settings.log_det_method != "cholesky"`` (the ``"slogdet"`` opt-in is untouched);
801+
- the factorization raises ``LinAlgError`` (``scipy.linalg.LinAlgError`` is
802+
``numpy.linalg.LinAlgError``), i.e. the matrix is not positive-definite;
803+
- the factor's diagonal is not finite (non-finite input).
804+
805+
The ``None`` fallbacks keep failure behaviour identical to before: an indefinite but
806+
nonsingular ``F + H`` is still solved by the LU and still makes the log determinant raise
807+
(or return 0.0 in test mode) through :meth:`_log_det_symmetric_from`; a singular one still
808+
makes the LU raise (or return its test-mode dummy).
809+
810+
``curvature_reg_matrix`` is a cached array that other properties read (the reduced matrix,
811+
the covariance), so it is copied, never factorized in place. The factor is one further
812+
``(n, n)`` array held for the inversion's lifetime.
813+
"""
814+
if (
815+
self._xp is not np
816+
or self.settings.use_positive_only_solver
817+
or self.settings.log_det_method != "cholesky"
818+
):
819+
return None
820+
821+
import scipy.linalg
822+
823+
try:
824+
cho_factor = scipy.linalg.cho_factor(
825+
np.asarray(self.curvature_reg_matrix),
826+
lower=True,
827+
overwrite_a=False,
828+
check_finite=False,
829+
)
830+
except np.linalg.LinAlgError:
831+
return None
832+
833+
if not np.all(np.isfinite(np.diag(cho_factor[0]))):
834+
return None
835+
836+
return cho_factor
837+
770838
@cached_property
771839
def reconstruction_reduced(self) -> np.ndarray:
772840
"""
@@ -1019,7 +1087,9 @@ def log_det_curvature_reg_matrix_term(self) -> float:
10191087
10201088
On the NumPy path with the positive-only solver and the default ``"cholesky"`` method this term is read
10211089
off the Cholesky factor the reconstruction's non-negative least squares solve has already built, rather
1022-
than factorizing ``F + reg_coeff*H`` a second time.
1090+
than factorizing ``F + reg_coeff*H`` a second time. On the NumPy path with the positive-negative solver it
1091+
is read off the Cholesky factor the reconstruction solved with (:attr:`_curvature_reg_matrix_cho_factor`),
1092+
for the same reason.
10231093
10241094
``fnnls_cholesky`` finishes holding a factor ``U`` of ``M[P][:, P]``, where ``M`` is
10251095
``curvature_reg_matrix_reduced`` and ``P`` is the solve's final passive set -- 1485 of the 1500 columns
@@ -1093,6 +1163,17 @@ def log_det_curvature_reg_matrix_term(self) -> float:
10931163
# not be answered with a fabricated 0.0 here.
10941164
pass
10951165

1166+
# Positive-negative NumPy path: read the log determinant off the Cholesky factor the
1167+
# reconstruction solves with (see `_curvature_reg_matrix_cho_factor`), rather than
1168+
# factorizing the same matrix again. Only when every linear object is regularized, so the
1169+
# reduced matrix is the full `curvature_reg_matrix` that was factorized. Reading the factor
1170+
# does not evaluate the reconstruction.
1171+
if self.all_linear_obj_have_regularization:
1172+
cho_factor = self._curvature_reg_matrix_cho_factor
1173+
1174+
if cho_factor is not None:
1175+
return 2.0 * np.sum(np.log(np.diag(cho_factor[0])))
1176+
10961177
return self._log_det_symmetric_from(self.curvature_reg_matrix_reduced)
10971178

10981179
@staticmethod
Lines changed: 280 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,280 @@
1+
"""
2+
One Cholesky factor of ``F + lambda*H`` shared by the positive-negative solve and its log determinant.
3+
4+
On the NumPy backend with ``use_positive_only_solver=False`` a likelihood call factorized the same
5+
symmetric positive-definite matrix twice: ``np.linalg.solve`` (an LU) for ``reconstruction`` and
6+
``np.linalg.cholesky`` for ``log_det_curvature_reg_matrix_term`` (PyAutoArray#633). Both now read
7+
``AbstractInversion._curvature_reg_matrix_cho_factor``.
8+
9+
These tests pin, against the two-factorization values computed here independently:
10+
11+
- the reconstruction (rtol 1e-10) and the log determinant (rtol 1e-9), on imaging and
12+
interferometer inversions, and that one factorization serves both;
13+
- the failure semantics: an indefinite-but-nonsingular ``F + H`` is still solved by the LU and still
14+
makes the log determinant raise (or return 0.0 in test mode); a singular one still makes the solve
15+
raise (or return its test-mode dummy);
16+
- the routes that must not use the factor: ``"slogdet"``, the positive-only solver, and a partially
17+
regularized system's log determinant.
18+
"""
19+
20+
import numpy as np
21+
import pytest
22+
import scipy.linalg
23+
24+
import autoarray as aa
25+
26+
from autoarray import fixtures
27+
28+
29+
def _two_factorization_values(inversion):
30+
"""The pre-#633 computation: an LU solve plus a separate dense Cholesky."""
31+
matrix = np.asarray(inversion.curvature_reg_matrix)
32+
33+
reconstruction = np.linalg.solve(matrix, np.asarray(inversion.data_vector))
34+
35+
reduced = np.asarray(inversion.curvature_reg_matrix_reduced)
36+
37+
log_det = 2.0 * np.sum(np.log(np.diag(np.linalg.cholesky(reduced))))
38+
39+
return reconstruction, log_det
40+
41+
42+
def _count_cho_factor(monkeypatch):
43+
original = scipy.linalg.cho_factor
44+
45+
calls = []
46+
47+
def _wrapped(*args, **kwargs):
48+
calls.append(1)
49+
return original(*args, **kwargs)
50+
51+
monkeypatch.setattr(scipy.linalg, "cho_factor", _wrapped)
52+
53+
return calls
54+
55+
56+
def _imaging_inversion(settings, linear_obj_list=None):
57+
return aa.Inversion(
58+
dataset=fixtures.make_masked_imaging_7x7(),
59+
linear_obj_list=linear_obj_list or [fixtures.make_rectangular_mapper_7x7_3x3()],
60+
settings=settings,
61+
)
62+
63+
64+
def _interferometer_inversion(settings):
65+
return aa.Inversion(
66+
dataset=fixtures.make_interferometer_7(),
67+
linear_obj_list=[fixtures.make_rectangular_mapper_7x7_3x3()],
68+
settings=settings,
69+
)
70+
71+
72+
def _with_matrix(inversion, matrix, data_vector):
73+
"""Seed the cached `curvature_reg_matrix` / `data_vector` with a chosen system."""
74+
inversion.__dict__["curvature_reg_matrix"] = matrix
75+
inversion.__dict__["data_vector"] = data_vector
76+
return inversion
77+
78+
79+
@pytest.mark.parametrize(
80+
"inversion_from",
81+
[_imaging_inversion, _interferometer_inversion],
82+
ids=["imaging", "interferometer"],
83+
)
84+
def test__shared_factor__matches_two_factorizations__one_factorization(
85+
monkeypatch, inversion_from
86+
):
87+
calls = _count_cho_factor(monkeypatch)
88+
89+
inversion = inversion_from(aa.Settings(use_positive_only_solver=False))
90+
91+
reconstruction_ref, log_det_ref = _two_factorization_values(inversion)
92+
93+
# log det first: it must read the factor without evaluating the reconstruction.
94+
log_det = inversion.log_det_curvature_reg_matrix_term
95+
96+
assert "reconstruction" not in inversion.__dict__
97+
98+
reconstruction = np.asarray(inversion.reconstruction)
99+
100+
assert len(calls) == 1
101+
assert inversion._curvature_reg_matrix_cho_factor is not None
102+
103+
assert reconstruction == pytest.approx(reconstruction_ref, rel=1.0e-10, abs=0.0)
104+
assert log_det == pytest.approx(log_det_ref, rel=1.0e-9)
105+
106+
# the factor is a copy: `F + H` and `F` are left intact for the reduced matrix / covariance
107+
assert np.asarray(inversion.curvature_reg_matrix) == pytest.approx(
108+
np.asarray(inversion.curvature_matrix)
109+
+ np.asarray(inversion.regularization_matrix),
110+
rel=1.0e-12,
111+
)
112+
113+
114+
def test__shared_factor__random_spd_system__matches_solve_and_dense_cholesky():
115+
rng = np.random.default_rng(3)
116+
117+
n = 9 # the rectangular 3x3 mapper's parameter count
118+
119+
a = rng.normal(size=(n, n))
120+
matrix = a @ a.T + 1.0e-3 * np.eye(n)
121+
data_vector = rng.normal(size=n)
122+
123+
inversion = _with_matrix(
124+
_imaging_inversion(aa.Settings(use_positive_only_solver=False)),
125+
matrix=matrix,
126+
data_vector=data_vector,
127+
)
128+
129+
assert np.asarray(inversion.reconstruction) == pytest.approx(
130+
np.linalg.solve(matrix, data_vector), rel=1.0e-10, abs=0.0
131+
)
132+
assert inversion.log_det_curvature_reg_matrix_term == pytest.approx(
133+
2.0 * np.sum(np.log(np.diag(np.linalg.cholesky(matrix)))), rel=1.0e-9
134+
)
135+
136+
137+
def _indefinite_nonsingular(n=9):
138+
matrix = np.eye(n)
139+
matrix[0, 0] = -1.0
140+
return matrix
141+
142+
143+
def test__indefinite_matrix__lu_solve_kept__log_det_raises(monkeypatch):
144+
monkeypatch.delenv("PYAUTO_TEST_MODE", raising=False)
145+
146+
matrix = _indefinite_nonsingular()
147+
data_vector = np.arange(1.0, 10.0)
148+
149+
inversion = _with_matrix(
150+
_imaging_inversion(aa.Settings(use_positive_only_solver=False)),
151+
matrix=matrix,
152+
data_vector=data_vector,
153+
)
154+
155+
assert inversion._curvature_reg_matrix_cho_factor is None
156+
157+
# as before: the LU solves an indefinite nonsingular system ...
158+
assert np.asarray(inversion.reconstruction) == pytest.approx(
159+
np.linalg.solve(matrix, data_vector), rel=1.0e-12
160+
)
161+
162+
# ... and the dense Cholesky log determinant rejects it
163+
with pytest.raises(np.linalg.LinAlgError):
164+
inversion.log_det_curvature_reg_matrix_term
165+
166+
167+
def test__indefinite_matrix__test_mode__log_det_guard_returns_zero(monkeypatch):
168+
monkeypatch.setenv("PYAUTO_TEST_MODE", "1")
169+
170+
inversion = _with_matrix(
171+
_imaging_inversion(aa.Settings(use_positive_only_solver=False)),
172+
matrix=_indefinite_nonsingular(),
173+
data_vector=np.ones(9),
174+
)
175+
176+
assert inversion.log_det_curvature_reg_matrix_term == 0.0
177+
178+
179+
def test__singular_matrix__solve_raises__test_mode_dummy(monkeypatch):
180+
matrix = np.ones((9, 9))
181+
182+
monkeypatch.delenv("PYAUTO_TEST_MODE", raising=False)
183+
184+
inversion = _with_matrix(
185+
_imaging_inversion(aa.Settings(use_positive_only_solver=False)),
186+
matrix=matrix,
187+
data_vector=np.arange(1.0, 10.0),
188+
)
189+
190+
with pytest.raises(np.linalg.LinAlgError):
191+
inversion.reconstruction
192+
193+
monkeypatch.setenv("PYAUTO_TEST_MODE", "1")
194+
195+
inversion = _with_matrix(
196+
_imaging_inversion(aa.Settings(use_positive_only_solver=False)),
197+
matrix=matrix,
198+
data_vector=np.arange(1.0, 10.0),
199+
)
200+
201+
assert (np.asarray(inversion.reconstruction) == 1.0).all()
202+
203+
204+
def test__non_finite_matrix__no_factor__unchanged_routes(monkeypatch):
205+
monkeypatch.delenv("PYAUTO_TEST_MODE", raising=False)
206+
207+
matrix = np.eye(9)
208+
matrix[4, 4] = np.nan
209+
210+
inversion = _with_matrix(
211+
_imaging_inversion(aa.Settings(use_positive_only_solver=False)),
212+
matrix=matrix,
213+
data_vector=np.ones(9),
214+
)
215+
216+
assert inversion._curvature_reg_matrix_cho_factor is None
217+
218+
219+
def test__slogdet_method__no_factor__unchanged(monkeypatch):
220+
calls = _count_cho_factor(monkeypatch)
221+
222+
inversion = _imaging_inversion(
223+
aa.Settings(use_positive_only_solver=False, log_det_method="slogdet")
224+
)
225+
226+
matrix = np.asarray(inversion.curvature_reg_matrix_reduced)
227+
228+
assert inversion.log_det_curvature_reg_matrix_term == pytest.approx(
229+
np.linalg.slogdet(matrix)[1], rel=1.0e-12
230+
)
231+
assert np.asarray(inversion.reconstruction) == pytest.approx(
232+
np.linalg.solve(matrix, np.asarray(inversion.data_vector)), rel=1.0e-12
233+
)
234+
assert inversion._curvature_reg_matrix_cho_factor is None
235+
assert calls == []
236+
237+
238+
def test__positive_only_solver__no_factor(monkeypatch):
239+
calls = _count_cho_factor(monkeypatch)
240+
241+
inversion = _imaging_inversion(aa.Settings(use_positive_only_solver=True))
242+
243+
_ = inversion.reconstruction
244+
_ = inversion.log_det_curvature_reg_matrix_term
245+
246+
assert inversion._curvature_reg_matrix_cho_factor is None
247+
assert calls == []
248+
249+
250+
def test__partially_regularized__solve_uses_factor__log_det_dense_reduced():
251+
dataset = fixtures.make_masked_imaging_7x7()
252+
253+
grid = aa.Grid2D.from_mask(mask=dataset.mask)
254+
255+
mapping_matrix = np.full(fill_value=0.5, shape=(9, 2))
256+
mapping_matrix[0, 0] = 0.8
257+
mapping_matrix[1, 1] = 0.4
258+
259+
linear_obj = aa.m.MockLinearObjFuncList(
260+
parameters=2, grid=grid, mapping_matrix=mapping_matrix
261+
)
262+
263+
inversion = _imaging_inversion(
264+
aa.Settings(use_positive_only_solver=False),
265+
linear_obj_list=[linear_obj, fixtures.make_rectangular_mapper_7x7_3x3()],
266+
)
267+
268+
assert not inversion.all_linear_obj_have_regularization
269+
270+
reconstruction_ref, log_det_ref = _two_factorization_values(inversion)
271+
272+
assert np.asarray(inversion.reconstruction) == pytest.approx(
273+
reconstruction_ref, rel=1.0e-10, abs=0.0
274+
)
275+
assert inversion._curvature_reg_matrix_cho_factor is not None
276+
277+
# the factor is of the full matrix; the log det is of the mapper-only reduced matrix
278+
assert inversion.log_det_curvature_reg_matrix_term == pytest.approx(
279+
log_det_ref, rel=1.0e-9
280+
)

0 commit comments

Comments
 (0)