Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 17 additions & 5 deletions scripts/imaging/pixelized/fixed_light_library.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand All @@ -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."
),
Expand Down Expand Up @@ -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))")
Expand Down
50 changes: 43 additions & 7 deletions scripts/misc/likelihood_breakdown/library_solver_injection.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
44 changes: 27 additions & 17 deletions scripts/misc/likelihood_breakdown/xla_attribution.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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).
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down
34 changes: 31 additions & 3 deletions scripts/misc/test/test_fixed_light_library.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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.

Expand Down
12 changes: 6 additions & 6 deletions scripts/misc/test/test_xla_attribution.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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"


Expand Down Expand Up @@ -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"
Expand Down
Loading