Skip to content

feat: Delaunay-pixelization JIT round-trip support in galaxy→inversion path#361

Merged
Jammy2211 merged 1 commit intomainfrom
feature/fit-imaging-pytree-delaunay
Apr 19, 2026
Merged

feat: Delaunay-pixelization JIT round-trip support in galaxy→inversion path#361
Jammy2211 merged 1 commit intomainfrom
feature/fit-imaging-pytree-delaunay

Conversation

@Jammy2211
Copy link
Copy Markdown
Collaborator

Summary

Enables jax.jit(analysis.fit_from) to succeed for models with a Delaunay
pixelization source + Hilbert image-mesh + AdaptSplit regularization. Without
this change the JIT path raises AttributeError: 'NoneType' object has no attribute 'array' inside GalaxiesToInversion.image_plane_mesh_grid_list, because the
Galaxy keys in adapt_images.galaxy_image_plane_mesh_grid_dict are the
original Python objects and the traced Galaxy is a fresh instance produced by
pytree unflatten — identity-keyed dict lookup fails.

API Changes

Behaviour-only change in GalaxiesToInversion.image_plane_mesh_grid_list: when
the galaxy-identity-keyed mesh-grid dict lookup returns None and the dict
contains exactly one entry, that entry is now used by insertion order.

No public signature changes. NumPy path is unaffected (it never hit the fallback).

See full details below.

Test Plan

  • pytest test_autogalaxy/ -q --ignore=test_autogalaxy/aggregator — 836 passed
  • scripts/jax_likelihood_functions/imaging/delaunay_pytree.py PASSes —
    NumPy and JIT log_likelihoods agree to rtol ~4e-15
  • Already-shipped pytree variants (lp, rectangular, mge_group) still pass

Upstream PR

None — this ships as a standalone library change.

Follow-up

The one-entry fallback is a deliberately narrow workaround. It is correct for the
current PoC (single pixelised source) but silently picks the wrong grid when there
are two or more pixelised sources. The principled fix — attach pytree_token to
Galaxy mirroring LightProfileLinear so identity survives pytree round-trips —
is filed at admin_jammy/prompt/autolens/galaxy_pytree_token.md in the workspace
repo and should land before any multi-source-pixelization JIT variant is added to
the queue (e.g. fit_imaging_pytree_delaunay_mge.md — item 10 in the queue).

Full API Changes (for automation & release notes)

Changed Behaviour

  • GalaxiesToInversion.image_plane_mesh_grid_list
    when adapt_images.galaxy_image_plane_mesh_grid_dict[galaxy] returns None
    (because the Galaxy instance has been reconstructed by a pytree unflatten
    cycle and no longer has identity-equality with the dict's keys), the method
    now falls back to the single value in the dict when the dict has length 1.
    Previously this case appended None to the result list, leading to a crash
    downstream when the mesh grid was accessed.

Rationale

  • Unblocks the Delaunay-pixelization PoC for jax.jit(analysis.fit_from).
  • NumPy path is unchanged — identity lookup succeeds without the fallback there.
  • Multi-pixelised-source case intentionally not fixed here; tracked by the
    galaxy_pytree_token.md follow-up prompt.

🤖 Generated with Claude Code

…r jax.jit

Adds a positional fallback in ``GalaxiesToInversion.image_plane_mesh_grid_list``
for the case where ``adapt_images.galaxy_image_plane_mesh_grid_dict[galaxy]``
returns None because the Galaxy instance is a fresh object produced by a
pytree unflatten cycle rather than the one stored as a dict key. When the
dict contains exactly one entry, that entry is returned by insertion order
— correct for the one-pixelised-source case (Delaunay / Hilbert image-mesh
fits) which is what the fit-imaging-pytree-delaunay PoC exercises.

This is a deliberately narrow workaround; the principled fix is to attach
a pytree_token to Galaxy (mirroring LightProfileLinear) so __hash__/__eq__
survive pytree round-trips. That fix is tracked separately in
admin_jammy/prompt/autolens/galaxy_pytree_token.md.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
@Jammy2211 Jammy2211 added the pending-release Merged PR awaiting inclusion in the next release build label Apr 19, 2026
@Jammy2211
Copy link
Copy Markdown
Collaborator Author

Workspace PR: PyAutoLabs/autolens_workspace_test#38

@Jammy2211 Jammy2211 merged commit 57a1291 into main Apr 19, 2026
5 checks passed
@Jammy2211 Jammy2211 deleted the feature/fit-imaging-pytree-delaunay branch April 19, 2026 20:53
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

pending-release Merged PR awaiting inclusion in the next release build

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant