From 2432bdc2e18fa8b6a65d98414dfdcb5350bd63ba Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 16 Oct 2026 11:32:55 +0100 Subject: [PATCH] fix: re-pin abstract.py anchors and probe the shared-Cholesky route after PyAutoArray#634 (#423) PyAutoArray#634 (02aa2e48) inserted 81 lines in inversion/abstract.py and routes the NumPy positivity-off solve through a shared Cholesky factor. Re-pin every xla_attribution abstract.py anchor (AST + difflib between 02aa2e48^ and 02aa2e48, anchored content identical) and make positive_negative_probe count both route-(c) paths (lu_solve, shared_cholesky; calls = their sum). Co-Authored-By: Claude Opus 5.5 --- .../imaging/pixelized/fixed_light_library.py | 22 ++++++-- .../library_solver_injection.py | 50 ++++++++++++++++--- .../likelihood_breakdown/xla_attribution.py | 44 +++++++++------- scripts/misc/test/test_fixed_light_library.py | 34 +++++++++++-- scripts/misc/test/test_xla_attribution.py | 12 ++--- 5 files changed, 124 insertions(+), 38 deletions(-) diff --git a/scripts/imaging/pixelized/fixed_light_library.py b/scripts/imaging/pixelized/fixed_light_library.py index 49071187a..036e5cf8e 100644 --- a/scripts/imaging/pixelized/fixed_light_library.py +++ b/scripts/imaging/pixelized/fixed_light_library.py @@ -702,7 +702,7 @@ def max_abs_rel(got, reference): ) #: Positivity off. ``Inversion.solve_ids_to_keep`` returns ``None`` the moment -#: this is false (abstract.py:540), so edge zeroing goes with it — which is why +#: this is false (abstract.py:568-569), so edge zeroing goes with it — which is why #: no route-(c) number is ever quoted without its Δevidence and its negatives. _settings_positive_negative = al.Settings( use_border_relocator=True, @@ -963,6 +963,8 @@ def _active_set_at(Q, q, fixed0, budget): _log_evidence_positive_negative = float(fit_positive_negative.figure_of_merit) _pn_calls = int(_pn_counts["calls"]) +_pn_lu_calls = int(_pn_counts["lu_solve"]) +_pn_cholesky_calls = int(_pn_counts["shared_cholesky"]) _pn_ids_to_keep = fit_positive_negative.inversion.solve_ids_to_keep # The phase-0 A2 row: one unconstrained Cholesky of the FULL S3 system. Route @@ -980,16 +982,25 @@ def _active_set_at(Q, q, fixed0, budget): _pn_negative = _x_positive_negative < 0.0 positive_negative_block = { "route_verified": ( - "reconstruction_positive_negative_from / xp.linalg.solve" + ( + "AbstractInversion.reconstruction / scipy.linalg.cho_solve " + "(shared Cholesky factor, PyAutoArray#634)" + if _pn_cholesky_calls > 0 and _pn_lu_calls == 0 + else "reconstruction_positive_negative_from / xp.linalg.solve" + if _pn_lu_calls > 0 and _pn_cholesky_calls == 0 + else "both: shared-Cholesky cho_solve and reconstruction_positive_negative_from" + ) if _pn_calls > 0 else "NOT VERIFIED — the positive-negative entry point was never called" ), "positive_negative_solver_calls": _pn_calls, + "positive_negative_lu_solve_calls": _pn_lu_calls, + "positive_negative_shared_cholesky_calls": _pn_cholesky_calls, "solve_ids_to_keep_is_none": _pn_ids_to_keep is None, "edge_zeroing_disabled": True, "edge_zeroing_note": ( "Inversion.solve_ids_to_keep returns None the moment use_positive_only_solver " - "is False (abstract.py:540), so turning positivity off silently turns edge " + "is False (abstract.py:568-569), so turning positivity off silently turns edge " "zeroing off with it. This route solves the FULL system; routes a/b/d/e solve " "the edge-zeroed one." ), @@ -1017,8 +1028,9 @@ def _active_set_at(Q, q, fixed0, budget): if _pn_calls == 0: raise AssertionError( - "use_positive_only_solver=False did not route through " - "reconstruction_positive_negative_from — route (c) is not what it claims to be." + "use_positive_only_solver=False did not route through the positive-negative " + "solve (shared-Cholesky cho_solve or reconstruction_positive_negative_from) — " + "route (c) is not what it claims to be." ) print(f" route verified: {positive_negative_block['route_verified']} ({_pn_calls} call(s))") diff --git a/scripts/misc/likelihood_breakdown/library_solver_injection.py b/scripts/misc/likelihood_breakdown/library_solver_injection.py index ca4de16c0..e33857be6 100644 --- a/scripts/misc/likelihood_breakdown/library_solver_injection.py +++ b/scripts/misc/likelihood_breakdown/library_solver_injection.py @@ -456,25 +456,61 @@ def positive_negative_probe(): """Count calls into the library's positive-**negative** solve, changing nothing. ``Settings(use_positive_only_solver=False)`` is supposed to route the - reconstruction to ``reconstruction_positive_negative_from`` — one - ``xp.linalg.solve`` — with edge zeroing silently off (``solve_ids_to_keep`` - returns ``None`` the moment positivity is off, abstract.py:540). This probe - is how the cell *verifies* that rather than asserting it in prose: the - counter must be non-zero on a route-(c) call and the positive-only counter - must be zero. + reconstruction to the unconstrained solve of ``F + λH``, with edge zeroing + silently off (``solve_ids_to_keep`` returns ``None`` the moment positivity is + off, abstract.py:568-569). This probe is how the cell *verifies* that rather + than asserting it in prose: the counter must be non-zero on a route-(c) call + and the positive-only counter must be zero. + + Since PyAutoArray#634 (main 02aa2e48) route (c) has two library paths, and + both are counted: + + - ``"shared_cholesky"`` — the NumPy path: ``AbstractInversion.reconstruction`` + returns ``scipy.linalg.cho_solve(self._curvature_reg_matrix_cho_factor, D)``, + the Cholesky factor the log determinant then reuses. Counted by wrapping + ``scipy.linalg.cho_solve`` for the duration of the context and counting + only calls made directly from ``reconstruction`` in + ``inversion/inversion/abstract.py`` — so the opt-in ``cho_solve`` + regularization term, the covariance matrix and anything else that calls + ``cho_solve`` never count. The factor is ``None`` whenever + ``use_positive_only_solver`` is on, and the positive-only branches return + before that line, so the positive-only path cannot trigger it. + - ``"lu_solve"`` — ``inversion_util.reconstruction_positive_negative_from`` + (one ``xp.linalg.solve``): the JAX path, and the NumPy fallback when the + factor is ``None`` (``log_det_method != "cholesky"``, or ``F + λH`` not + positive-definite). + + ``"calls"`` is their sum — the total number of route-(c) solves — so a + non-zero ``"calls"`` keeps meaning "route (c) ran", whichever path it took. """ + import sys + + import scipy.linalg from autoarray.inversion.inversion import inversion_util original = inversion_util.reconstruction_positive_negative_from - counts = {"calls": 0} + original_cho_solve = scipy.linalg.cho_solve + counts = {"calls": 0, "lu_solve": 0, "shared_cholesky": 0} def patched(data_vector, curvature_reg_matrix, xp=np): counts["calls"] += 1 + counts["lu_solve"] += 1 return original(data_vector=data_vector, curvature_reg_matrix=curvature_reg_matrix, xp=xp) + def cho_solve_observed(*args, **kwargs): + caller = sys._getframe(1).f_code + if caller.co_name == "reconstruction" and caller.co_filename.replace("\\", "/").endswith( + "inversion/inversion/abstract.py" + ): + counts["calls"] += 1 + counts["shared_cholesky"] += 1 + return original_cho_solve(*args, **kwargs) + _assert_signature_covers(original, patched) inversion_util.reconstruction_positive_negative_from = patched + scipy.linalg.cho_solve = cho_solve_observed try: yield counts finally: inversion_util.reconstruction_positive_negative_from = original + scipy.linalg.cho_solve = original_cho_solve diff --git a/scripts/misc/likelihood_breakdown/xla_attribution.py b/scripts/misc/likelihood_breakdown/xla_attribution.py index 93b71aa36..92113c9a5 100644 --- a/scripts/misc/likelihood_breakdown/xla_attribution.py +++ b/scripts/misc/likelihood_breakdown/xla_attribution.py @@ -62,9 +62,9 @@ This matters because the stages are not separable by the innermost frame alone. Both Bayesian-evidence log determinants are the *same* three lines of -``inversion/inversion/abstract.py`` (``_log_det_symmetric_from``, 969-1013 at d4298445) and are -told apart only by their caller — ``log_det_curvature_reg_matrix_term`` (1096) or -``log_det_regularization_matrix_term`` (1175). So a :class:`StageRule` may declare +``inversion/inversion/abstract.py`` (``_log_det_symmetric_from``, 1037-1081 at 02aa2e48) and are +told apart only by their caller — ``log_det_curvature_reg_matrix_term`` (1177) or +``log_det_regularization_matrix_term`` (1256). So a :class:`StageRule` may declare a ``requires`` frame that must also appear somewhere in the same stack. Fusions carry their constituents @@ -617,14 +617,15 @@ def satisfied_by(self, frames: Sequence[Frame]) -> bool: #: 2026-09-25 after PyAutoArray#571 moved everything below line 602; re-pinned #: 2026-09-30 after PyAutoArray#586 moved everything below line 305 by +9; re-pinned #: 2026-09-30 again after PyAutoArray#591 inserted one import line at 10 (+1 below it) -#: and eight lines in ``data_subtracted_dict`` at old 880 (+9 below it)). -LIBRARY_ANCHOR_REVISION = "PyAutoArray main d4298445 (2026-09-30)" +#: and eight lines in ``data_subtracted_dict`` at old 880 (+9 below it); re-pinned +#: 2026-10-16 after PyAutoArray#634 inserted 81 lines at 763/770/1022/1094 (+12, +68, +70, +81)). +LIBRARY_ANCHOR_REVISION = "PyAutoArray main 02aa2e48 (2026-10-15)" #: ``_log_det_symmetric_from`` — the shared body of both log determinants. -_ABSTRACT_LOG_DET_BODY = (969, 1013) +_ABSTRACT_LOG_DET_BODY = (1037, 1081) #: ``log_det_curvature_reg_matrix_term`` (decorator to last line). -_ABSTRACT_LOG_DET_CURVATURE_CALLER = (1015, 1096) +_ABSTRACT_LOG_DET_CURVATURE_CALLER = (1083, 1177) #: ``log_det_regularization_matrix_term`` (decorator to last line). -_ABSTRACT_LOG_DET_REGULARIZATION_CALLER = (1116, 1177) +_ABSTRACT_LOG_DET_REGULARIZATION_CALLER = (1197, 1258) #: ``curvature_reg_matrix``'s single-line ``self._xp.add(F, H)``. _ABSTRACT_CURVATURE_ADD_LINE = 399 #: ``curvature_reg_matrix_reduced`` (the ``mapper_indices`` gathers). @@ -659,12 +660,21 @@ def satisfied_by(self, frames: Sequence[Frame]) -> bool: #: ``data_subtracted_dict`` at old line 880 (every range below it a further +8, so #: +9; ``mapped_reconstruction``'s 831-893, which ends in that function, became #: 832-902). Every span re-derived from the AST of both revisions; content identical. +#: RE-PINNED 2026-10-16 against PyAutoArray ``main`` 02aa2e48: #634 added the shared +#: Cholesky ``cho_solve`` route inside ``reconstruction`` (+12 at old 763), the new +#: ``_curvature_reg_matrix_cho_factor`` before ``reconstruction_reduced`` (+56 at old 770), +#: two docstring lines in ``log_det_curvature_reg_matrix_term`` (old 1022) and its +#: shared-factor branch (+11 at old 1094). Every range below 763 is unchanged; every +#: range below 1096 shifted +81. The NEW lines are NumPy-only (the factor is ``None`` on +#: JAX), so they land inside the spans that enclose them (``reconstruction_scatter`` +#: 714-830 -> 714-898, the curvature caller 1015-1096 -> 1083-1177) and never trace. +#: Every span re-derived from the AST of both revisions; anchored lines identical. #: ``LIBRARY_ANCHOR_REVISION`` names that revision in the provenance. STAGE_MAP: tuple[StageRule, ...] = ( # --- the two log determinants: same lines, told apart by their caller ---- - # ``_log_det_symmetric_from`` 969-1013; its callers - # ``log_det_curvature_reg_matrix_term`` 1015-1096 and - # ``log_det_regularization_matrix_term`` 1116-1177. + # ``_log_det_symmetric_from`` 1037-1081; its callers + # ``log_det_curvature_reg_matrix_term`` 1083-1177 and + # ``log_det_regularization_matrix_term`` 1197-1258. StageRule( "log_det_curvature_reg", _ARRAY + _INVERSION_ABSTRACT, @@ -696,17 +706,17 @@ def satisfied_by(self, frames: Sequence[Frame]) -> bool: # --- the F + lambda*H add, and the two [ids][:, ids] gathers ------------- StageRule("curvature_reg_add", _ARRAY + _INVERSION_ABSTRACT, (378, 399)), StageRule("curvature_reg_reduce_gather", _ARRAY + _INVERSION_ABSTRACT, (401, 425)), - # ``reconstruction`` spans 669-768. 706-713 is the edge-zeroed branch + # ``reconstruction`` spans 669-780. 706-713 is the edge-zeroed branch # (``ids_to_keep is not None``): the two ``[ids_to_keep]`` subsets. 714-744 # is the partial solve plus the scatter back to full shape; 746-762 is the # else branch, which is what the Delaunay family runs (no edge zeroing). - # 714-830 runs on through ``reconstruction_reduced`` / ``_dict`` / - # ``source_quantity_dict_from``, as the 2026-09-16 map's 615-721 did. + # 714-898 runs on through ``_curvature_reg_matrix_cho_factor`` (NumPy-only) / + # ``reconstruction_reduced`` / ``_dict`` / ``source_quantity_dict_from``, as the 2026-09-16 map's 615-721 did. StageRule("edge_subset_gather", _ARRAY + _INVERSION_ABSTRACT, _ABSTRACT_EDGE_SUBSET), - StageRule("regularization_term", _ARRAY + _INVERSION_ABSTRACT, (903, 967)), - StageRule("mapped_reconstruction", _ARRAY + _INVERSION_ABSTRACT, (832, 902)), + StageRule("regularization_term", _ARRAY + _INVERSION_ABSTRACT, (971, 1035)), + StageRule("mapped_reconstruction", _ARRAY + _INVERSION_ABSTRACT, (900, 970)), StageRule("reconstruction_scatter", _ARRAY + _INVERSION_ABSTRACT, (669, 705)), - StageRule("reconstruction_scatter", _ARRAY + _INVERSION_ABSTRACT, (714, 830)), + StageRule("reconstruction_scatter", _ARRAY + _INVERSION_ABSTRACT, (714, 898)), # --- the PDIP / active-set solve ---------------------------------------- # ``inversion_util.py`` ranges (re-anchored 2026-09-24, 681938ae; re-pinned # 2026-09-25, 5f8a8dee, after #571): 12-93 diff --git a/scripts/misc/test/test_fixed_light_library.py b/scripts/misc/test/test_fixed_light_library.py index 61f335da2..32bafeebf 100644 --- a/scripts/misc/test/test_fixed_light_library.py +++ b/scripts/misc/test/test_fixed_light_library.py @@ -357,9 +357,8 @@ def test_positivity_off_routes_to_the_positive_negative_solver(tiny_fit): ) reconstruction = np.asarray(fit.inversion.reconstruction, dtype=float) - assert counts["calls"] >= 1, ( - "positivity off did not reach reconstruction_positive_negative_from" - ) + assert counts["calls"] >= 1, "positivity off did not reach the positive-negative solve" + assert counts["calls"] == counts["lu_solve"] + counts["shared_cholesky"] assert fit.inversion.solve_ids_to_keep is None, "edge zeroing survived positivity being off" expected = np.linalg.solve( @@ -372,6 +371,35 @@ def test_positivity_off_routes_to_the_positive_negative_solver(tiny_fit): ) +def test_positive_negative_probe_is_silent_on_the_positive_only_route(tiny_fit): + """The probe's control: positivity ON must count zero route-(c) solves. + + The shared-Cholesky counter wraps ``scipy.linalg.cho_solve``; this pins that + the positive-only path (whose factor gate is closed) never trips it, and that + the probe restores the library entry points on exit. + """ + import scipy.linalg + from autoarray.inversion.inversion import inversion_util + + al, dataset, tracer = tiny_fit + cho_solve = scipy.linalg.cho_solve + lu_solve = inversion_util.reconstruction_positive_negative_from + + with lsi.positive_negative_probe() as counts: + fit = al.FitImaging( + dataset=dataset, + tracer=tracer, + settings=al.Settings(use_border_relocator=True), + xp=np, + ) + np.asarray(fit.inversion.reconstruction, dtype=float) + float(fit.figure_of_merit) + + assert counts == {"calls": 0, "lu_solve": 0, "shared_cholesky": 0} + assert scipy.linalg.cho_solve is cho_solve + assert inversion_util.reconstruction_positive_negative_from is lu_solve + + def test_positivity_on_keeps_edge_zeroing(tiny_fit): """The control: with positivity on, the solve IS subset to solve_ids_to_keep. diff --git a/scripts/misc/test/test_xla_attribution.py b/scripts/misc/test/test_xla_attribution.py index d03288531..18038e059 100644 --- a/scripts/misc/test/test_xla_attribution.py +++ b/scripts/misc/test/test_xla_attribution.py @@ -129,9 +129,9 @@ def _encode_stack_frame_index( ] _LOCATIONS = [ (1, 1, 399, 15), # 1 abstract.py:399 curvature_reg_matrix (F + lambda*H) - (1, 2, 977, 20), # 2 abstract.py:977 _log_det_symmetric_from (the cholesky) - (1, 3, 1069, 15), # 3 abstract.py:1069 log_det_curvature_reg_matrix_term - (1, 4, 1148, 15), # 4 abstract.py:1148 log_det_regularization_matrix_term + (1, 2, 1045, 20), # 2 abstract.py:1045 _log_det_symmetric_from (the cholesky) + (1, 3, 1139, 15), # 3 abstract.py:1139 log_det_curvature_reg_matrix_term + (1, 4, 1229, 15), # 4 abstract.py:1229 log_det_regularization_matrix_term (2, 5, 586, 21), # 5 convolver.py:586 the shared rfft2 (4, 6, 136, 20), # 6 imaging/abstract.py:136 operated_mapping_matrix_list (3, 7, 295, 12), # 7 delaunay.py:295 the visibility walk while_loop @@ -225,9 +225,9 @@ def test_frames_are_innermost_first_and_walk_the_parent_chain(stack_index): """Frame 4 is the Cholesky *called by* log_det_curvature_reg_matrix_term.""" frames = stack_index.frames_for(4) assert len(frames) == 2 - assert frames[0].line == 977 + assert frames[0].line == 1045 assert frames[0].function == "_log_det_symmetric_from" - assert frames[1].line == 1069 + assert frames[1].line == 1139 assert frames[1].function == "log_det_curvature_reg_matrix_term" @@ -470,7 +470,7 @@ def test_the_harness_logdet_candidate_splits_into_solver_and_log_det_rows(): dense_body = xa.Frame( file=_LIB + "inversion/inversion/abstract.py", function="_log_det_symmetric_from", - line=977, + line=1045, ) patched = xa.Frame(file=_HARNESS_LOGDET, function="patched_log_det", line=400) assert xa.stage_for_frames((dense_body, patched)) == "log_det_curvature_reg"