Skip to content

feat: SimulatorImaging.via_galaxies_from auto-default xp from parent use_jax#442

Merged
Jammy2211 merged 1 commit into
mainfrom
feature/simulator-imaging-use-jax
May 24, 2026
Merged

feat: SimulatorImaging.via_galaxies_from auto-default xp from parent use_jax#442
Jammy2211 merged 1 commit into
mainfrom
feature/simulator-imaging-use-jax

Conversation

@Jammy2211
Copy link
Copy Markdown
Collaborator

Summary

Companion PR to the PyAutoArray PR adding use_jax=True to aa.SimulatorImaging. The autogalaxy via_galaxies_from override gains an xp=None parameter and defaults to self._xp from the parent — same pattern as the PyAutoLens via_tracer_from companion.

After all 3 PRs merge:

simulator = ag.SimulatorImaging(exposure_time=300.0, psf=psf, use_jax=True)
dataset = simulator.via_galaxies_from(galaxies=galaxies, grid=grid)  # JAX-backed

API Changes

  • Changed signature: ag.SimulatorImaging.via_galaxies_from(galaxies, grid, xp=None)xp now defaults to None. New behaviour: falls back to self._xp (inherited from aa.SimulatorImaging).

See full details below.

Test Plan

  • PyAutoGalaxy 918/918 tests pass.
  • Eager cross-xp parity (via autolens_workspace_test parity script, bundled in the workspace_test follow-up).
Full API Changes

Changed signature

  • ag.SimulatorImaging.via_galaxies_from(galaxies, grid, xp=None) — gained xp parameter (previously not accepted). New behaviour: xp=None falls back to self._xp.

Migration

  • No breaking changes. Existing callers without xp= continue to work.

Depends on

  • Jammy2211/PyAutoArray PR adding use_jax to aa.SimulatorImaging. Merge PyAutoArray first.

🤖 Generated with Claude Code

… use_jax

Changes via_galaxies_from to accept xp=None and fall back to self._xp
inherited from aa.SimulatorImaging. When the simulator is constructed with
use_jax=True, via_galaxies_from now routes xp=jnp through
galaxies.padded_image_2d_from and via_image_from without the user passing
xp explicitly.

Depends on Jammy2211/PyAutoArray PR adding the use_jax constructor flag +
_xp property to the parent aa.SimulatorImaging.

Part of Phase 2 PR 2 of z_features/jax_user_intro.md.

Issue: PyAutoArray#334

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@Jammy2211 Jammy2211 added the pending-release PR queued for the next release build label May 24, 2026
@Jammy2211 Jammy2211 merged commit 86defc1 into main May 24, 2026
4 of 6 checks passed
@Jammy2211 Jammy2211 deleted the feature/simulator-imaging-use-jax branch May 24, 2026 15:40
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.

1 participant