Skip to content

feat: scripts/jax_grad/interferometer/ — lp.py + mge.py from scratch (task 7/9)#31

Merged
Jammy2211 merged 1 commit intomainfrom
feature/autogalaxy-wst-jax-grad-interferometer
May 6, 2026
Merged

feat: scripts/jax_grad/interferometer/ — lp.py + mge.py from scratch (task 7/9)#31
Jammy2211 merged 1 commit intomainfrom
feature/autogalaxy-wst-jax-grad-interferometer

Conversation

@Jammy2211
Copy link
Copy Markdown
Collaborator

Summary

Create scripts/jax_grad/interferometer/{lp.py, mge.py} from scratch — autolens has no interferometer jax_grad reference today. Each script runs jax.value_and_grad on the autogalaxy AnalysisInterferometer likelihood path and asserts the gradient is finite, the shape matches model.total_free_parameters, and the gradient is not all-zero.

This is task 7/9 of the autogalaxy_workspace_test parity epic (#5).

Scripts Changed

  • scripts/jax_grad/interferometer/__init__.py — new (empty namespace marker).
  • scripts/jax_grad/interferometer/lp.py — new. Single galaxy with a Sersic bulge on dataset/interferometer/jax_test with a (256,256) pixel_scales=0.1 radius=3.0 real-space mask and TransformerDFT (mirrors jax_likelihood_functions/interferometer/lp.py). Body uses jax.value_and_grad(fitness.call) and the four canonical assertions. Local run: 6.9s, gradient shape (7,).
  • scripts/jax_grad/interferometer/mge.py — new. Same dataset/mask/transformer; model swapped to a single-basis MGE (mge_model_from, 20 gaussians, centre_prior_is_uniform=True) on a single galaxy. Same body and assertions. Local run: 12.5s, gradient shape (4,).
  • smoke_tests.txt — appended both new scripts immediately after jax_grad/imaging/mge.py.

Notes

Layout divergence from autolens (continued). Autolens has flat top-level jax_grad/imaging_*.py and no interferometer jax_grad scripts at all. This PR continues the subfolder convention established in PR #29 (task 6) — jax_grad/{imaging,interferometer,multi}/ on autogalaxy. The retrofit question for autolens was raised in PR #29; suggest filing the follow-up after task 8 lands so the migration covers all three subdirectories at once.

lp.Sersic vs lp_linear.Sersic. Task 6 (jax_grad/imaging/lp.py) used lp_linear.Sersic. This task uses plain lp.Sersic to match the validated jax_likelihood_functions/interferometer/lp.py pattern. The interferometer + JAX inversion path differs from imaging; sticking with the proven setup.

No env_vars.yaml change. The jax_grad/ substring override added in PR #29 (unset PYAUTO_SMALL_DATASETS + PYAUTO_DISABLE_JAX) already matches jax_grad/interferometer/.

Test Plan

  • Both new scripts pass locally via .github/scripts/run_smoke.py with the env override applied:
    • jax_grad/interferometer/lp.py — exit 0 in 6.9s, grad shape (7,)
    • jax_grad/interferometer/mge.py — exit 0 in 12.5s, grad shape (4,)

🤖 Generated with Claude Code

Closes #30.

…(task 7/9)

Create autogalaxy interferometer jax_grad scripts under a new
jax_grad/interferometer/ subfolder. Autolens has no interferometer
jax_grad reference today, so these are greenfield ports modelled on
the matching jax_likelihood_functions/interferometer/{lp,mge}.py
scripts.

Each script wraps the AnalysisInterferometer log-likelihood in
jax.value_and_grad and asserts the gradient is finite, has the
expected free-parameter shape, and is not all-zero.

The jax_grad/ env_vars override added in task 6 (PR #29) already
covers jax_grad/interferometer/ — no env_vars.yaml change needed.

Closes #30.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

pending-release PR queued for the next release build

Projects

None yet

Development

Successfully merging this pull request may close these issues.

feat: scripts/jax_grad/interferometer/ — lp.py + mge.py from scratch (task 7/9)

1 participant