perf: slim LOS tests and remove JAX solver test

[Jammy2211]

Overview

The PyAutoLens test suite takes ~188s. The 5 LOSSampler tests consume ~107s (57%) and the 3 NegativeKappa tests add ~15s — all for an in-development LOS feature. The aggregator tests add ~20s of database/fixture overhead better suited as integration tests. The test_real_example_jax solver test adds 12s of JIT overhead now covered by smoke/integration tests. Moving these out and keeping minimal unit coverage should cut ~139s (~74% reduction).

Plan

• Move the full test_los.py test suite to autolens_workspace_developer (in-development feature)
• Keep one consolidated LOS unit test in PyAutoLens that exercises LOSSampler end-to-end
• Move all aggregator tests (test_aggregator_fit_imaging.py, test_aggregator_fit_interferometer.py, test_aggregator_tracer.py, conftest.py) to autolens_workspace_test/scripts/ as integration tests
• Remove test_real_example_jax from test_solver.py — JAX coverage is handled by smoke/integration tests
• Verify all remaining PyAutoLens tests pass and compare before/after runtime

Detailed implementation plan

Affected Repositories

• PyAutoLens (primary) — remove/slim tests
• autolens_workspace_developer — receive moved LOS tests
• autolens_workspace_test — receive moved aggregator tests

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoLens	main	1 modified (docs)
./autolens_workspace_developer	main	1 modified (root.log)
./autolens_workspace_test	main	TBD

Suggested branch: feature/unit-test-profiling

Implementation Steps

1. Move LOS tests — Copy test_autolens/lens/test_los.py to autolens_workspace_developer/los/test_los.py, preserving all existing tests and imports

2. Consolidate LOS in PyAutoLens — Replace TestLOSSampler (5 tests, ~107s) and TestNegativeKappa (3 tests, ~15s) with a single consolidated test that:

   • Creates a LOSSampler with small params (2 planes, m_max=1e9, cone_radius_arcsec=1.0)
   • Calls galaxies_from() and asserts: returns galaxies, has NFWTruncatedSph halos, has MassSheet entries with negative kappa, all redshifts match plane centres, sheet count matches plane count
   • Keep the pure-math tests (TestComovingDistance, TestComovingVolume, TestLosPlanes, TestMassRatio, TestNumberOfHalos, TestSampleHaloMasses, TestSamplePositions, TestSampleConcentrations) — these are fast (<0.01s each)

3. Move aggregator tests — Move test_autolens/aggregator/ contents to autolens_workspace_test/scripts/aggregator/:

   • test_aggregator_fit_imaging.py (~9s, 3 tests)
   • test_aggregator_fit_interferometer.py (~8s, 3 tests)
   • test_aggregator_tracer.py (~5s, 2 tests)
   • conftest.py (shared fixtures)
   • Remove the test_autolens/aggregator/ directory from PyAutoLens

4. Remove JAX solver test — Delete test_real_example_jax from test_autolens/point/triangles/test_solver.py (~12s saved)

5. Verify — Run full PyAutoLens suite, confirm pass, target ~50s total (from 188s)

Key Files

• test_autolens/lens/test_los.py — slim down LOSSampler/NegativeKappa tests
• test_autolens/aggregator/ — move entirely to autolens_workspace_test
• test_autolens/point/triangles/test_solver.py — remove test_real_example_jax
• autolens_workspace_developer/los/test_los.py — new home for full LOS test suite
• autolens_workspace_test/scripts/aggregator/ — new home for aggregator integration tests

Original Prompt

Click to expand starting prompt

Profile and speed up unit tests in @PyAutoLens.

Goal

Reduce the test suite runtime by identifying and optimizing slow tests. The source code is stable, so safe changes to test parameters are acceptable — but prompt the user before changing any numerical assertion values.

Approach (proven on PyAutoFit)

1. Profile: Run pytest <test_dir> --durations=0 -q to get wall-clock times for all tests
2. Identify: Find the top 30+ slowest tests and read their source to understand bottlenecks
3. Categorize: Group by cause — large data arrays, excessive iterations/max_steps, expensive fixtures, real sampler initialization (Dynesty)
4. Optimize (in order of impact):
   • Reduce n_obs, data array sizes, range(N) loop counts where tests use more data than needed for their assertion tolerances
   • Reduce max_steps, n_iters, maxcall values in EP/sampler tests that validate plumbing not convergence
   • Reduce number_of_steps in grid search fixtures to shrink the job count (e.g. 10->5 gives 25 jobs instead of 100)
   • Adjust query/assertion values if data range changes (e.g. if range(100)->range(50), update assertions that referenced value 50 to value 25)
   • Leave tests with inherent sampler initialization overhead (DynestyStatic with maxcall<=5) as-is
5. Verify: Run full suite, confirm all tests pass, compare before/after total time

What NOT to change

• Don't change numerical reference values in assertions without user approval
• Don't change fixture scoping to module/session — this conflicts with function-scoped set_config_path autouse fixtures
• Don't mock real sampler tests — they test actual sampling behavior
• Don't change test_messages.py::test_normal_simplex — the slowness is numerical quadrature integration

Known pattern: FactorValue hash collision

If you encounter a KeyError: <GaussianPrior id=0> in AbstractJacobian.grad(), this is a known bug where FactorValue (id=0) hash-collides with a GaussianPrior that gets id=0. This was fixed in PyAutoFit PR #1182 — check that the fix is present in the installed autofit version.

Target

25-35% reduction in total test runtime. On PyAutoFit this took 61s -> 41s.

[Jammy2211]

Library PR Created — 2026-04-06

PR

• perf: reduce test suite from 188s to 22s (88% faster) #421

Status

Workspace changes needed. Next: /start_workspace

API Changes

None — internal test changes only. LOS tests, aggregator tests, and JAX solver test removed from unit suite. Need to be re-added as integration tests in autolens_workspace_test and autolens_workspace_developer.

[Jammy2211]

Shipped — 2026-04-06

PRs

• perf: reduce test suite from 188s to 22s (88% faster) #421
• feat: add aggregator integration test scripts autolens_workspace_test#12
• https://github.com/PyAutoLabs/autolens_workspace_developer/pull/1

Summary

Reduced PyAutoLens unit test suite from 188s (286 tests) to 22s (245 tests) — 88% faster. Moved LOS tests to autolens_workspace_developer (in-development feature), aggregator tests to autolens_workspace_test as standalone scripts, removed redundant JAX solver test. Additionally optimised test harness by mocking plt.tight_layout and plt.colorbar in plot_patch fixture.

Session Notes

• Profiling revealed plt.tight_layout() accounted for ~50% of plotter test time (font metric computation for 12-panel subplots). Mocking it in the test fixture was a zero-risk fix.
• plt.colorbar was the next biggest contributor (~58% of per-panel rendering). Using MagicMock() as the return value handled all downstream attribute access.
• The LOS LOSSampler tests are inherently slow (~13s even with minimal params) due to nested scipy.integrate.quad calls in negative_kappa_from. Reducing planes from 2+2 to 1+1 cut this to ~8s.
• autolens_workspace_developer had no git remote — set up https://github.com/PyAutoLabs/autolens_workspace_developer during this session.