feat: autolens interferometer JAX visualization scripts#87
Merged
Conversation
added 2 commits
May 8, 2026 21:44
Closes #86. - scripts/interferometer/visualization_jax.py (NEW) — exercises VisualizerInterferometer.visualize on the JIT-cached fit_for_visualization path. - scripts/interferometer/modeling_visualization_jit.py (NEW) — caching probe + live Nautilus run with linear MGE profiles, asserts the jit-cache fires during quick updates and fit.png lands on disk. - config/build/env_vars.yaml — split the interferometer/visualization override into NumPy-only and JAX-only entries; add interferometer/modeling_visualization_jit override. Mirrors the imaging analogues. Phase 1A of PyAutoPrompt/issued/jax_visualization.md. Depends on PyAutoLabs/PyAutoLens#500.
Without this cleanup the script would pass on first run but fail on re-run because Nautilus resumes from cached samples.csv, skips live sampling entirely, and the quick-update visualizer never fires — so analysis_mge2._jitted_fit_from is never set and the assertion at the bottom of Part 2 raises AttributeError. Fix: rmtree output/<output_root>/<name>/ in addition to the existing scripts/.../images/ rmtree, forcing every run to actually exercise the live sampling path.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds the missing
interferometer/visualization_jax.pyandinterferometer/modeling_visualization_jit.pyscripts in autolens_workspace_test, mirroring the imaging analogues. Phase 1A ofPyAutoPrompt/issued/jax_visualization.md(the JAX visualization roadmap). Together with the upstream PyAutoLens fix (#500), this closes the autolens imaging + interferometer coverage for the JIT-cached visualization path.Closes #86.
Scripts Changed
scripts/interferometer/visualization_jax.py(NEW) — exercisesVisualizerInterferometer.visualizeon the JIT-cachedfit_for_visualizationpath. Usesenable_pytrees()+register_model(model)from the start, notry/exceptwrapper around the visualize call.scripts/interferometer/modeling_visualization_jit.py(NEW) — two-part: Part 1 caching probe (callsanalysis.fit_for_visualizationtwice and asserts the second call is significantly faster than the first); Part 2 live Nautilus run with linear MGE profiles (asserts_jitted_fit_fromis cached on the analysis instance andfit.pnglands on disk during quick updates). Cleansoutput/<path_prefix>/<name>/before launching Nautilus so the search always runs fresh — without this, reruns silently resume from cachedsamples.csvand the JIT cache is never populated, breaking the assertion.config/build/env_vars.yaml— split the existinginterferometer/visualizationoverride into NumPy-only (interferometer/visualization.py) and JAX-only (interferometer/visualization_jax) entries, mirroring the autolens imaging split shipped in fix: register_model in visualization_jax.py to actually exercise JIT path #85. Add aninterferometer/modeling_visualization_jitoverride mirroring the imaging analogue.Upstream PR
PyAutoLabs/PyAutoLens#500 —
AnalysisInterferometer.__init__kwargs passthrough. Required for these scripts to constructAnalysisInterferometer(use_jax_for_visualization=True, ...)withoutTypeError. Library-first merge gate enforces this PR merges only after #500.Test Plan
python scripts/interferometer/visualization_jax.py(with JAX enabled) —PILOT SUCCEEDED — JAX-backed interferometer visualization produced fit.png/tracer.png.python scripts/interferometer/modeling_visualization_jit.py(with JAX enabled) — JIT cache fires during Nautilus,fit.pngproduced, no KeyError fromlinear_light_profile_intensity_dict.🤖 Generated with Claude Code