feat: autolens point_source JAX visualization scripts#91
Merged
Conversation
Closes #90. - scripts/point_source/visualization.py (NEW) — NumPy baseline calling VisualizerPoint.visualize on the NumPy fall-through path. - scripts/point_source/visualization_jax.py (NEW) — JIT-cached fit_for_visualization via use_jax_for_visualization=True, with enable_pytrees() + register_model(model) from the start. - scripts/point_source/modeling_visualization_jit.py (NEW) — caching probe + live Nautilus quick-update run with image-plane chi-squared (FitPositionsImagePairAll). Cleans output/<path_prefix>/<name>/ before launching Nautilus so reruns don't silently skip live sampling. - config/build/env_vars.yaml — point_source/visualization_jax + point_source/modeling_visualization_jit overrides. Phase 1B of PyAutoPrompt/issued/jax_visualization.md. Image-plane only (source-plane JIT still blocked per scripts/CLAUDE.md L132). Depends on PyAutoLabs/PyAutoLens#506.
1 task
This was referenced May 14, 2026
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
point_source/visualization.py,point_source/visualization_jax.py, andpoint_source/modeling_visualization_jit.pyscripts in autolens_workspace_test. Phase 1B ofPyAutoPrompt/issued/jax_visualization.md— point-source was the only autolens dataset type with zero visualization coverage of any flavour. Together with the upstream PyAutoLens fix (#506), this closes the autolens imaging + interferometer + point_source coverage for the JIT-cached visualization path.Closes #90.
Scripts Changed
scripts/point_source/visualization.py(NEW) — NumPy baseline. Builds an Isothermal+PointFlux model and callsVisualizerPoint.visualizedirectly withuse_jax=False. Assertsfit.pnglands on disk.scripts/point_source/visualization_jax.py(NEW) — JIT-cached path. Addsenable_pytrees()+register_model(model)and constructsal.AnalysisPoint(use_jax=True, use_jax_for_visualization=True, fit_positions_cls=al.FitPositionsImagePairAll, ...). No try/except wrapper. Assertsfit.pnglands.scripts/point_source/modeling_visualization_jit.py(NEW) — two-part:analysis.fit_for_visualization(instance)twice, asserts second is significantly faster, asserts_jitted_fit_from is not None).n_live=50, n_like_max=1500, iterations_per_quick_update=500and image-plane chi-squared. Cleans bothscripts/point_source/images/modeling_visualization_jit/ANDoutput/scripts/point_source/images/modeling_visualization_jit/point_image_plane/before the Nautilus call so reruns don't silently resume from cachedsamples.csv(lesson from PR feat: autolens interferometer JAX visualization scripts #87).config/build/env_vars.yaml— addpoint_source/visualization_jax(unsetPYAUTO_DISABLE_JAX, PYAUTO_SMALL_DATASETS) andpoint_source/modeling_visualization_jit(unsetPYAUTO_DISABLE_JAX, PYAUTO_SMALL_DATASETS, PYAUTO_TEST_MODE, PYAUTO_FAST_PLOTS) overrides, mirroring the imaging and interferometer analogues.All three scripts use
FitPositionsImagePairAll(image-plane chi-squared) only. Source-plane chi-squared (FitPositionsSource) is still JIT-blocked perscripts/CLAUDE.mdL132 — out of scope for this PR; will need its own follow-up when the blocker lifts.The model deliberately omits a free
cosmologyparameter — cosmology distance calc caches global state and breaks JIT round-trip (per the existingjax_likelihood_functions/point_source/image_plane.pyL144-147 caveat).Upstream PR
PyAutoLabs/PyAutoLens#506 —
AnalysisPoint.__init__**kwargspassthrough. Required for these scripts to constructAnalysisPoint(use_jax_for_visualization=True, ...)withoutTypeError. Library-first merge gate enforces this PR merges only after #506.Test Plan
python scripts/point_source/visualization.py—NumPy point-source visualization produced fit.png.JAX_ENABLE_X64=True python scripts/point_source/visualization_jax.py—PILOT SUCCEEDED — JAX-backed point-source visualization produced fit.png.JAX_ENABLE_X64=True python scripts/point_source/modeling_visualization_jit.py— JIT cache fires during Nautilus,fit.png files produced: 1.🤖 Generated with Claude Code