refactor: defer heavy imports — import autolens 4.1s → ~1.9s

[Jammy2211]

Overview

import autolens costs ~4.1 s (full example-script stanza ~4.4 s), and profiling shows ~2.5 s of it is heavy packages loaded eagerly but not needed at import: the jax chain (nufftax + blackjax + optax, 1.75 s — 43%, measured by control test), IPython, sqlalchemy, numba (decoration-time), and astropy. All 457 autolens_workspace scripts pay this per-process (~30 min of pure import per full serial smoke sweep), and smoke runs with JAX disabled, so its jax import there is pure waste. This task defers each of these to first use across PyAutoFit, PyAutoArray and PyAutoNerves, targeting import autolens ≤ ~1.9 s with zero public-API change.

Plan

• Defer the jax chain: lazy af.NSS (PEP 562) in PyAutoFit; lazy nufftax load + batcher patch in PyAutoArray's TransformerNUFFT.
• Defer IPython, sqlalchemy (lazy sa proxy + annotation fixes), and the eager autofit.database pulls in PyAutoFit.
• Defer numba to first call via a lazy-compile wrapper in autoarray.numba_util.jit (fixes all 29 decorated functions centrally).
• Defer astropy in autonerves.fitsable; dedupe the workspace-version warning (currently fires 3× per process).
• Verify: import-time absence proof, full suites in all three repos, NSS / NUFFT / aggregator behaviour runs, targeted workspace smoke incl. # ENV: jax scripts.
• Scoped out: matplotlib deferral (enters via the autolens/autogalaxy visualizer chain; every workspace script imports aplt anyway, so no smoke/user win — optional follow-up for bare-library import only).

Detailed implementation plan

Affected Repositories

• PyAutoFit (primary)
• PyAutoArray
• PyAutoNerves

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoFit	main	clean
./PyAutoArray	main	clean
./PyAutoNerves	main	clean

Suggested branch: feature/lazy-heavy-imports

Note: PyAutoFit is also claimed by stored-sample-reconstruction-guard (#1486, library half shipped as PR#1504) — deliberate file-disjoint override, recorded in Mind active.md. The version-stamp-sync-guards branch also touches autofit/__init__.py (version-stamp lines); coordinate at merge if both are in flight.

Implementation Steps (PyAutoFit, ~0.9 s)

1. autofit/__init__.py — remove eager imports of NSS (L98), GridSearchAggregator (L15), Aggregator (L37), Query (L42); add PEP 562 module __getattr__ resolving these four lazily, caching in globals() (pattern: autolens/__init__.py:165-181). NSS stays fully supported — lazy, never removed.
2. autofit/non_linear/fitness.py — move from IPython.display import clear_output (L3) into the single use site (~L439, quick-update branch).
3. autofit/database/sqlalchemy_.py — make sa/declarative lazy proxies: __getattr__ imports sqlalchemy on first attribute access and caches; fall back to the existing MockSQlAlchemy on ImportError.
4. Add from __future__ import annotations to modules evaluating sa.orm.Session in signatures at def time (Py3.12): non_linear/settings.py, search/mcmc/abstract_mcmc.py, search/nest/abstract_nest.py, search/nest/nss/search.py, search/nest/nautilus/search.py, plus a full per-file sweep of every sqlalchemy_ importer. Verify no eager __annotations__ introspection in these modules.
5. autofit/non_linear/paths/database.py — defer Fit and Info imports to function-local at their use sites (~L112, L236-243, L345).
6. autofit/non_linear/grid/grid_search/result_builder.py — from __future__ import annotations + move from autofit.database import Prior (annotation-only) under if TYPE_CHECKING.

Implementation Steps (PyAutoArray, ~0.75 s)

7. autoarray/operators/transformer.py — replace module-level nufftax try-import (L26-29) and version-gated _patch_nufftax_batchers() call (L90-94) with an idempotent _load_nufftax() (import, 0.6.x check, batcher patch — a global JAX registration that must precede the first traced call) invoked at the top of TransformerNUFFT.__init__ (existing _nufftax is None guard at ~L646). All other _nufftax uses are in TransformerNUFFT methods.
8. autoarray/numba_util.py — jit() returns a wrapper that on first call imports numba (falling back to the plain function on ModuleNotFoundError), compiles with the same nopython/cache/parallel/fastmath options, caches, delegates. Sweep for attribute access on decorated functions (.py_func etc.) first.

Implementation Steps (PyAutoNerves, ~0.17 s + noise)

9. autonerves/fitsable.py — remove module-level astropy try-import (L12-15, keep TYPE_CHECKING block); lazy _fits() accessor used at the six runtime sites (~L63, 86, 88, 90, 210, 238).
10. autonerves/workspace.py — module-level _warned_messages set; the three warnings.warn sites (~L227, 237, 257) skip byte-identical repeats, so import autolens warns once, not three times. The direct-call tests use distinct tmp_path roots — unaffected.

Key Files

• PyAutoFit/autofit/__init__.py — lazy NSS/Aggregator/Query/GridSearchAggregator
• PyAutoFit/autofit/database/sqlalchemy_.py — lazy sa proxy
• PyAutoArray/autoarray/operators/transformer.py — lazy nufftax + batcher patch
• PyAutoArray/autoarray/numba_util.py — lazy-compile jit wrapper
• PyAutoNerves/autonerves/fitsable.py, PyAutoNerves/autonerves/workspace.py

Verification

• Timing: 3× cold import autolens before/after; target ≤ ~1.9 s (matplotlib retained by design).
• Absence proof: python -X importtime -c "import autolens" shows no jax/jaxlib/blackjax/optax/nufftax/IPython/sqlalchemy/numba/astropy.
• Behaviour: af.NSS/af.Aggregator/af.Query/af.GridSearchAggregator resolve; run autofit_workspace_test/scripts/searches/NSS.py, an interferometer TransformerNUFFT script, a database/aggregator flow.
• Full pytest suites in PyAutoFit, PyAutoArray, PyAutoNerves; downstream autogalaxy/autolens suites via ship_library; unittest-nojax leg stays green.
• Targeted workspace smoke incl. # ENV: jax scripts before PRs open.

Follow-up (not this task)

• Brain RefactorDecision flags pyautonerves as "unwitnessed" — its witness map lacks test_autonerves; fix the map.
• Optional: matplotlib deferral for bare-library import (unit-test startup), via the autolens/autogalaxy visualizer chain.

Original Prompt

Click to expand starting prompt

Can you do a profiling of how long import times on autolens_workspace
example scripts and performance in general is, and assess if we can speed
it up as it is important for running smoke tests and general user
performance being fast

Full measured profile in PyAutoMind/active/import_time_lazy_heavy_imports.md.

[Jammy2211]

Verification results (worktree feature/lazy-heavy-imports)

Headline: import autolens 4.07 s → ~2.1 s (full script stanza ~4.4 s → ~2.0 s). python -X importtime proves jax, jaxlib, blackjax, optax, nufftax, IPython, sqlalchemy, numba and astropy are all absent from the bare import tree. The autonerves workspace-version warning now fires once per process (was 3×).

Unit suites (all green)

Repo	Result
PyAutoNerves	157 passed
PyAutoArray	1062 passed
PyAutoFit	2005 passed, 16 skipped (pre-existing NSS skips: local blackjax 1.5 has no blackjax.ns; _HAS_NSS False on canonical main too)
PyAutoGalaxy	1111 passed
PyAutoLens	538 passed

Workspace smoke (pyauto-heart smoke --root <worktree>, all six workspaces)

Workspace	Passed	Failed	Total
autofit_workspace	10	0	10
autogalaxy_workspace	16	0	16
autolens_workspace	37	0	37
autolens_workspace_test	19	4 (+1 timeout)	24
euclid pipeline	6	0	6
HowToLens	6	0	6

The 4 autolens_workspace_test failures are pre-existing, not from this branch — a control run against canonical main (same harness, same env) fails the identical four scripts:

• imaging/jax_likelihood/lp.py — vmap result is ~half the hardcoded pin; matches the in-flight positions-LH penalty-accumulation fix (orphan worktree commit fb1aefe0b "sum PositionsLH penalties instead of doubling the last entry" — the pin was calibrated against the doubling behaviour).
• interferometer/jax_likelihood/rectangular.py — misses its pin by 0.0124% against rtol 1e-4 (stale-pin drift).
• interferometer/jax_likelihood/mge.py, multi_dataset/jax_likelihood/mge.py — same pin-drift family.

The imaging/subhalo_recovery.py timeout (300 s cap) occurred only under the six-workspace sweep's load: an uncontended worktree rerun passes it at 252.1 s (canonical: 252.5 s — no slowdown from this branch) and reproduces exactly the canonical result, 20/24 with the same four pre-existing failures.

Behaviour checks

• af.NSS / af.Aggregator / af.Query / af.GridSearchAggregator / af.db resolve via PEP 562 and construct.
• In-memory db.open_database("sqlite://") round-trip works (sqlalchemy loads on first use).
• TransformerNUFFT numpy/jax visibilities agree exactly; transform_mapping_matrix and image_from work; a pickled instance in a state-reset module (multiprocessing-worker simulation) works.
• numba lazy materialization compiles nopython cross-calls correctly (module globals rebound to CPUDispatchers on first call).
• End-to-end Drawer model fit runs through the annotated search plumbing.
• No from autofit import * users exist; sphinx autosummary resolves lazy names via getattr.

[Jammy2211]

Library PRs created (all pending-release, all gates green)

Repo	Tests	PR
PyAutoFit	2005 passed, 16 skipped (pre-existing NSS env skips)	#1506
PyAutoArray	1062 passed	PyAutoLabs/PyAutoArray#451
PyAutoNerves	157 passed	PyAutoLabs/PyAutoNerves#151
PyAutoGalaxy	1111 passed	PyAutoLabs/PyAutoGalaxy#575
PyAutoLens	538 passed	PyAutoLabs/PyAutoLens#705

Heart readiness at ship: GREEN, score 100. Six-workspace smoke against the branch: identical to canonical main (the four autolens_workspace_test jax_likelihood failures are control-proven pre-existing — see the verification comment above).

Result: import autolens 4.07 s → ~2.1 s; full example-script stanza ~4.4 s → ~2.0 s; jax/jaxlib/blackjax/optax/nufftax/IPython/sqlalchemy/numba/astropy all absent from the bare import tree; workspace-version warning deduped to once per process.

Workspace impact: None — all five PRs declare "None — internal changes only"; no workspace migration or new demos needed (option iii). Merge is the remaining human decision.

Follow-ups noted (not this task): Brain witness map lacks test_autonerves; optional matplotlib deferral for bare-library import; the pre-existing jax_likelihood pin failures belong to the positions-LH penalty-accumulation arc + stale pins.

[Jammy2211]

Shipped — all five PRs merged 2026-08-19

PyAutoFit#1506, PyAutoArray#451, PyAutoNerves#151, PyAutoGalaxy#575, PyAutoLens#705 — all 17 CI checks green (incl. unittest-nojax), remote branches deleted.

Final measured result (warm cache; earlier ~2.1 s readings were polluted by the fresh worktree's one-time matplotlib font-cache build and first-run bytecode compilation):

• import autolens: 4.07 s → 1.2–1.3 s (~3.2× faster) — beats the ≤1.9 s acceptance target.
• Import tree self-time: 3.22 s / 3494 modules → 1.27 s / 1814 modules.
• Remaining big items, all deliberate: scipy 0.31 s (pynufft's eager import + scattered uses), matplotlib stack 0.33 s (scoped out — every workspace script imports aplt anyway), the PyAuto libraries' own module tree ~0.38 s.

Closing. Follow-ups tracked as Mind drafts: pre-existing autolens_workspace_test jax_likelihood pin failures (positions-LH penalty arc + stale pins), Brain witness map missing test_autonerves.