test: add log_likelihood_function and Fitness.call_wrap assertions to _test workspaces

[Jammy2211]

Overview

PR #504 in PyAutoLens fixed a long-standing bug where the CPU branch of AnalysisImaging.log_likelihood_function returned fit.log_likelihood instead of fit.figure_of_merit. For pixelization-inversion fits these differ by the regularization log-det terms of the Bayesian log evidence, so every CPU-driven nested-sampler search with a pixelization source was silently optimising the wrong objective — nested samplers drifted to outer_coefficient ≈ 0 (noise-overfit degenerate mode) instead of the physical Bayesian maximum.

The bug went unnoticed because the existing unit test used a purely parametric Sersic model where figure_of_merit == log_likelihood. PR #504 adds one focused unit-level regression. This task extends the _test workspaces with broader, end-to-end "likelihood sanity" assertions so any future regression of this kind is caught early on realistic configurations, not just the minimal 7×7 unit-test fixture.

Plan

• Insert a __Likelihood Sanity__ block before the search starts in each _test script that fits a pixelization source. Build the prior-median instance, call analysis.log_likelihood_function(instance=instance), reconstruct the equivalent FitImaging / FitInterferometer from analysis.tracer_via_instance_from(...) + analysis.adapt_images_via_instance_from(...), and assert it equals fit.figure_of_merit while differing from fit.log_likelihood (so the test cannot pass tautologically).
• Add an equivalent Fitness.call_wrap assertion in the same block, using model.physical_values_from_prior_medians so the Nautilus-facing surface is also guarded.
• Cover both backends where the script exercises both — run the assertions once with use_jax=False and once with use_jax=True (only on scripts that already build a JAX analysis).
• Mirror the additions in interferometer scripts (AnalysisInterferometer + Fitness). No bug exists there today, but the guard prevents the same asymmetry from being introduced.
• _test workspaces only — do not touch user-facing autolens_workspace / autogalaxy_workspace tutorials.

Detailed implementation plan

Affected Repositories

• autolens_workspace_test (primary)
• autogalaxy_workspace_test

Work Classification

Workspace.

Branch Survey

Repository	Current Branch	Dirty?
./autolens_workspace_test	main	clean
./autogalaxy_workspace_test	main	clean

Suggested branch: feature/likelihood-function-assertions
Worktree root: ~/Code/PyAutoLabs-wt/likelihood-function-assertions/ (created later by /start_workspace)

Implementation Steps

1. Pre-run exploration (start of /start_workspace). Run each candidate script in its current form (with the sampler limits already in place) to confirm where the __Likelihood Sanity__ block fits naturally. Confirm insertion points with the user before editing.

2. Define a reusable assertion block. Inline (no shared module), so each _test script remains self-contained:

   """
   __Likelihood Sanity__
   
   Guard against regressions like PyAutoLens PR #504, where the CPU branch of
   ``log_likelihood_function`` silently returned ``fit.log_likelihood`` instead
   of ``fit.figure_of_merit``. For a pixelization source these differ by the
   regularization log-det terms of the Bayesian log evidence.
   """
   import pytest
   from autofit.non_linear.fitness import Fitness
   
   instance = model.instance_from_prior_medians()
   analysis_value = analysis.log_likelihood_function(instance=instance)
   
   tracer = analysis.tracer_via_instance_from(instance=instance)
   fit_adapt_images = analysis.adapt_images_via_instance_from(
       instance=instance, galaxies=tracer.galaxies
   )
   fit = al.FitImaging(
       dataset=dataset, tracer=tracer, adapt_images=fit_adapt_images,
       settings_inversion=analysis.settings_inversion,
   )
   
   assert analysis_value == pytest.approx(fit.figure_of_merit)
   assert fit.figure_of_merit != pytest.approx(fit.log_likelihood, rel=1e-6)
   
   fitness = Fitness(
       model=model, analysis=analysis, paths=None,
       fom_is_log_likelihood=True, resample_figure_of_merit=-1.0e99,
   )
   parameter_vector = model.physical_values_from_prior_medians
   assert fitness.call_wrap(parameter_vector) == pytest.approx(fit.figure_of_merit)

   The exact FitImaging / FitInterferometer constructor arguments may need adjusting per script (positions_likelihood, settings_inversion, transformer_class for interferometer). Mirror what the corresponding analysis.fit_from(...) builds internally.

3. Cover both backends where applicable. For scripts that already construct an analysis with use_jax=True (the modeling_visualization_jit*.py family), repeat the assertions once with the existing use_jax=True analysis and once with a sibling use_jax=False analysis built from the same dataset/model so both branches are protected.

4. Imaging — autolens_workspace_test:

   • scripts/imaging/model_fit.py — Isothermal + Delaunay pixelization (CPU/Nautilus). Add NumPy-path assertions.
   • scripts/imaging/modeling_visualization_jit_delaunay.py — PowerLaw + Delaunay (JAX). Add JAX-path + NumPy-path assertions. Reuse the existing instance_probe-style prior-median build.
   • scripts/imaging/modeling_visualization_jit_rectangular.py — sibling rectangular mesh (JAX). Same treatment.
   • scripts/imaging/modeling_visualization_jit.py — MGE source: confirm whether it includes a pixelization in any branch; if not, skip per the prompt's "fit a pixelization source" criterion (or add only the FoM==log_likelihood equality without the != guard).

5. Imaging — autogalaxy_workspace_test:

   • scripts/imaging/model_fit.py — inspect; if it fits a pixelization, add NumPy assertions for AnalysisImaging + Fitness.
   • scripts/imaging/modeling_visualization_jit.py — sibling JAX pixelization script. Same treatment.

6. Interferometer — autolens_workspace_test:

   • scripts/interferometer/model_fit.py — Isothermal + Delaunay pixelization. Use FitInterferometer and AnalysisInterferometer. Same assertion pattern adapted for the interferometer surface.
   • scripts/interferometer/modeling_visualization_jit.py — JAX interferometer with pixelization (if applicable). Mirror.

7. Interferometer — autogalaxy_workspace_test:

   • scripts/interferometer/modeling_visualization_jit.py — JAX interferometer (if applicable).

8. Run-time impact. Each __Likelihood Sanity__ block adds one extra likelihood evaluation per script plus one Fitness.call_wrap call (and at most one extra JIT compile on the JAX branch). For the JIT scripts this should be a small fraction of the existing sampler budget; verify by timing one before/after on modeling_visualization_jit_delaunay.py.

9. Validate. Run every modified script to confirm:

   • The __Likelihood Sanity__ block exits cleanly with all assertions passing.
   • The downstream search still completes within its existing n_like_max / iterations_per_quick_update budget.
   • On JAX scripts, no second compile is forced on fit_for_visualization after the sanity check (re-use the cached jitted callable where possible).

Key Files

• autolens_workspace_test/scripts/imaging/model_fit.py — NumPy-path pixelization assertions
• autolens_workspace_test/scripts/imaging/modeling_visualization_jit_delaunay.py — JAX + NumPy assertions
• autolens_workspace_test/scripts/imaging/modeling_visualization_jit_rectangular.py — JAX + NumPy assertions
• autolens_workspace_test/scripts/imaging/modeling_visualization_jit.py — assertions iff a pixelization branch is exercised
• autolens_workspace_test/scripts/interferometer/model_fit.py — AnalysisInterferometer + FitInterferometer
• autolens_workspace_test/scripts/interferometer/modeling_visualization_jit.py — interferometer JAX assertions (if pixelization)
• autogalaxy_workspace_test/scripts/imaging/model_fit.py — confirm pixelization, then NumPy assertions
• autogalaxy_workspace_test/scripts/imaging/modeling_visualization_jit.py — JAX assertions
• autogalaxy_workspace_test/scripts/interferometer/modeling_visualization_jit.py — JAX interferometer assertions (if applicable)

References

• PyAutoLens/test_autolens/imaging/model/test_analysis_imaging.py — see test__log_likelihood_function__returns_figure_of_merit_for_pixelization (lines 47–71) for the unit-level assertion pattern this scales up.
• PyAutoLens/autolens/imaging/model/analysis.py — the fixed log_likelihood_function; the assertions here guard against future edits that reintroduce the asymmetry.
• PyAutoFit/autofit/non_linear/fitness.py:209 — Fitness.call_wrap signature and behaviour.

Original Prompt

Click to expand starting prompt

We just fixed a long-standing bug in
@PyAutoLens/autolens/imaging/model/analysis.py where the CPU branch of
AnalysisImaging.log_likelihood_function returned fit.log_likelihood
instead of fit.figure_of_merit (PR #504, merged). For pixelization
inversion fits these differ by the regularization log-det terms of the
Bayesian log evidence, so every CPU-driven nested-sampler search with a
pixelization source was silently optimising the wrong objective —
nested samplers drifted to outer_coefficient ≈ 0 (noise-overfit
degenerate mode) instead of converging to the physical Bayesian
maximum.

The bug went unnoticed for a long time because the existing unit test
in @PyAutoLens/test_autolens/imaging/model/test_analysis_imaging.py
(test__figure_of_merit__matches_correct_fit_given_galaxy_profiles)
used a purely parametric Sersic model — where
figure_of_merit == log_likelihood — and therefore didn't exercise the
diverging branch. PR #504 adds one focused regression test for the
pixelization case.

We need to extend the _test workspaces with broader, end-to-end
assertions so that any future regression of this kind gets caught early
on realistic configurations, not just the minimal 7x7 unit-test fixture.

What needs adding:

1. Identify the regression-style integration scripts in
   @autolens_workspace_test/scripts/imaging/ (likely model_fit.py,
   modeling_visualization_jit_delaunay.py,
   modeling_visualization_jit_rectangular.py and similar) and in
   @autogalaxy_workspace_test/scripts/imaging/ that fit a pixelization
   source. For each, add a small block (does NOT need its own script —
   can sit inline before the search starts) that:

   • Builds the model's prior-median instance via
     instance = model.instance_from_prior_medians().
   • Computes analysis_value = analysis.log_likelihood_function(instance=instance).
   • Reconstructs the equivalent fit = al.FitImaging(dataset=..., tracer=..., adapt_images=..., settings=...)
     directly (via analysis.tracer_via_instance_from(instance) and
     analysis.adapt_images_via_instance_from(instance, galaxies=tracer.galaxies)).
   • Asserts analysis_value == pytest.approx(fit.figure_of_merit).
   • Asserts fit.figure_of_merit != pytest.approx(fit.log_likelihood, rel=1e-6)
     so the test cannot pass tautologically on a pixelization-less model.

2. Add an equivalent Fitness assertion. Build the Fitness wrapper the
   way Nautilus would:

   from autofit.non_linear.fitness import Fitness
   fitness = Fitness(
       model=model, analysis=analysis, paths=None,
       fom_is_log_likelihood=True, resample_figure_of_merit=-1.0e99,
   )
   parameter_vector = model.physical_values_from_prior_medians
   assert fitness.call_wrap(parameter_vector) == pytest.approx(fit.figure_of_merit)

   This guards the Nautilus-facing surface specifically. The
   single-instance log_likelihood_function check guards the
   AnalysisImaging surface; Fitness.call_wrap guards the conversion
   from parameter vector → instance → log-evidence that Nautilus
   actually invokes per sample.

3. Repeat the same two assertions for both backends if the script
   exercises both — i.e. compute the assertions once with
   use_jax=False and once with use_jax=True (where applicable on
   the workspace's CI runner). This makes the regression cover any
   future drift between the CPU and JAX branches.

4. Update @autolens_workspace_test/scripts/interferometer/ and
   @autogalaxy_workspace_test/scripts/interferometer/ analogously for
   AnalysisInterferometer + Fitness. The interferometer analysis
   already returns figure_of_merit correctly (no bug there today) but
   we want a guard in case the same asymmetry gets introduced.

5. Do NOT touch the autolens_workspace or autogalaxy_workspace user-
   facing tutorials. These assertions are integration-test
   infrastructure and belong in the _test workspaces only.

References (read first):

• @PyAutoLens/test_autolens/imaging/model/test_analysis_imaging.py — see
  the new unit test added in PR #504 (lines around
  test__log_likelihood_function__returns_figure_of_merit_for_pixelization)
  for the assertion pattern to mirror at integration scale.
• @PyAutoLens/autolens/imaging/model/analysis.py — the fixed
  log_likelihood_function; the assertions here are testing that
  future edits to this method don't reintroduce the asymmetry.
• @PyAutoFit/autofit/non_linear/fitness.py — Fitness.call_wrap
  signature and behaviour.

Before starting, run @autolens_workspace_test/scripts/imaging/model_fit.py
and one of the visualization_jit scripts in their current form to see
what they produce. Then propose where the assertion blocks should sit
in each file (before the search? in a dedicated __Likelihood Sanity__
section?), confirm with me, then implement.

[Jammy2211]

Shipped — 2026-05-18

PRs

• test: add __Likelihood Sanity__ blocks to pixelization scripts #103 (MERGED)
• test: add __Likelihood Sanity__ block to imaging/visualization.py autogalaxy_workspace_test#50 (MERGED)

Summary

Added inline __Likelihood Sanity__ regression-guard blocks before the Nautilus search in every integration script that fits a pixelization source. Each block asserts analysis.log_likelihood_function == fit.figure_of_merit, fit.figure_of_merit != fit.log_likelihood (kills tautological pass), and Fitness.call_wrap == fit.figure_of_merit — guarding the AnalysisImaging / AnalysisInterferometer surface AND the Nautilus-facing parameter-vector → instance → figure-of-merit conversion. JIT scripts cover both CPU + JAX backends; non-JIT scripts cover CPU only.

Session Notes

• Final scope is narrower than the original prompt. The audit found:
  • All modeling_visualization_jit.py (singular) scripts use MGE sources, no pixelization → out of scope (would have to drop the != guard to be non-tautological).
  • autogalaxy_workspace_test/scripts/imaging/model_fit.py uses parametric Sersic only → out of scope.
  • autogalaxy_workspace_test/scripts/interferometer/* has no pixelization script at all → out of scope. The autogalaxy AnalysisInterferometer + pixelization surface remains unexercised in integration tests; add coverage there only when a pixelization-source interferometer script is added for another reason. Library unit tests are still the only sentinel.
  • autogalaxy_workspace_test/scripts/imaging/visualization.py is the natural home for the autogalaxy guard — it already builds rectangular + Delaunay pixelization analyses with instance_from_prior_medians, so the block drops in inside the existing per-source loop on if has_inversion.
• Reconstruction simplified. Rather than rebuilding FitImaging / FitInterferometer manually with dataset_model, settings, etc. (the prompt's example was incomplete), the block just calls analysis.fit_from(instance) — the same call log_likelihood_function makes internally. The bug detection is unaffected (the bug was which attribute gets pulled off the fit, not how the fit is built).
• Penalty term sidestepped. A dedicated sanity_analysis_cpu is built without positions_likelihood_list so there's no log_likelihood_penalty to subtract from analysis_value. Keeps the assertions clean across all scripts.
• Local JAX validation was env-blocked on this laptop's PyAutoGPU venv — analysis.log_likelihood_function on use_jax=True hits RuntimeError: Unknown backend cuda. Available backends are ['cpu'] regardless of JAX_PLATFORMS=cpu. The CPU branch was fully validated; the JAX branch will be exercised on CI build.
• Smoke env validated locally for autogalaxy_workspace_test/scripts/imaging/visualization.py (the only modified file in either repo's smoke_tests.txt) using its config/build/env_vars.yaml override (unsets PYAUTO_FAST_PLOTS and PYAUTO_SMALL_DATASETS). Both pixelization iterations PASS.
• Future work: if the autogalaxy interferometer + pixelization surface ever gets an integration script, drop the same block in. Helper is identical — only ag.AnalysisInterferometer swaps in for ag.AnalysisImaging.