Multi-band grid rotation: cleaning up Qiuhan's prototype

[Jammy2211]

Tagging @qiuhan06 — this issue tracks finishing the multi-band grid-rotation work you prototyped on the dev_Q branches, so Hannah Stacey's JWST multi-band fits can model both pointing offsets and roll-angle misalignments through a single DatasetModel parameter.

How this started

From Slack earlier today (paraphrased):

Jam: Hannah mentioned modelling rotations in multi-band JWST fits — when you get a chance, could you push your PR so I can implement it properly? Doesn't need to be anything proper, with the right code Claude can do the rest.

Qiuhan: My previous implementation (transferring and rotating the grids) failed — source reconstruction behaved very weird and I thought I was missing something. I'll push what I have to dev_Q branches on PyAutoArray / PyAutoGalaxy / PyAutoLens.

Qiuhan: Worry is, I don't have a good description of the issue I see 😅.

Jam: Claude will figure it out from the code.

Diagnosis

Reproduced your failure mode on ephemeral worktrees patched with the autoarray-side Track 1 changes only. Five chi² ≈ 0 self-consistency tests covering offset, rotation, both, parametric source, and pixelization source — all pass.

What was actually going wrong: on main, AdaptImages.galaxy_name_image_plane_mesh_grid_dict was never being shifted by grid_offset, let alone rotated. For multi-band fits with an adaptive pixelization (Hilbert / KMeans), the cached mesh grid sits in the original image-plane frame while the data grid (correctly) sits in the offset-and-rotated frame — they don't share a coordinate system, so the source reconstruction is genuinely misaligned. Your adapt_images.py patch ("saving before updates" commit) already fixes this; you just hadn't re-tested it.

So the prototype was the right approach. We've forward-ported the minimum set of changes from dev_Q onto current main, dropped the *MultiBand mass profile workaround (PowerLawMultiBand, ExternalShearMultiBand, NFWMCRScatterLudlowSphMultiBand, PowerLawMultipoleMultiBand) and the unrelated mge_model_V2_from, cleaned up debug prints, and added regression tests.

What ships

• aa.DatasetModel.grid_rotation_angle — new degrees-CCW float parameter, default 0.0, mirroring the existing grid_offset pattern exactly.
• Grid2D.subtracted_and_rotated_from(offset, angle, xp) and Grid2DIrregular.subtracted_and_rotated_from(...) — shift-then-rotate transform, used by FitDataset.grids and AdaptImages.updated_via_instance_from.
• AdaptImages.updated_via_instance_from now takes a dataset_model argument and applies the same shift+rotate to cached mesh grids. This is the actual fix for the "weird source reconstruction".
• Library tests (1,917 passing): new subtracted_and_rotated_from unit tests in PyAutoArray + new chi² ≈ 0 simulate-and-fit tests in both test_autogalaxy/imaging/test_simulate_and_fit_imaging.py and test_autolens/imaging/test_simulate_and_fit_imaging.py for offset-only, rotation-only, and combined cases. These will catch any future regression of the latent mesh-grid misalignment bug.
• Workspace example: autolens_workspace/scripts/multi/features/dataset_offsets/{simulator,modeling}.py extended to inject and recover a 1.5° relative roll angle alongside the existing half-pixel offset.

PRs

Library-first ordering. Each repo's PR references this issue.

• PyAutoArray: feat(DatasetModel): add grid_rotation_angle for multi-band rotation PyAutoArray#312
• PyAutoGalaxy: feat(AdaptImages): rotate cached mesh grid with DatasetModel transforms PyAutoGalaxy#416
• PyAutoLens: feat(AnalysisImaging): plumb dataset_model into adapt_images_via_instance_from #512
• autolens_workspace: docs(multi/dataset_offsets): demonstrate grid_rotation_angle autolens_workspace#154
• autolens_workspace_test: test(multi): DatasetModel shift+rotation parity scripts autolens_workspace_test#97 — the definitive Delaunay parity regression guard

Verification

test_autoarray/ ............ 767 passed
test_autogalaxy/ ........... 873 passed
test_autolens/ ............. 277 passed

Plus the four self-consistent simulate-and-fit chi² ≈ 0 tests (three in test_autogalaxy, three in test_autolens) — each one simulates a profile with an offset centre and/or rotated ellipse, fits with an identity-frame profile + matching DatasetModel parameters, and asserts chi_squared = pytest.approx(0.0, abs=1e-4). These are the regression guards for this entire feature.

Time-to-fix narrative

For Qiuhan & anyone curious how this kind of work is shaped now: this whole job — diagnose the failure, port the prototype forward across three libraries + workspace, add ten regression tests, run all 1,917 unit tests green — was one CLI conversation. No JIRA tickets, no follow-up meetings, no "I'll get to it next sprint". The hardest part was the Phase 0 reproduction step that proved the math was right before changing any production code; everything after that was mechanical.

Closes once the four PRs merge.

[Jammy2211]

@qiuhan06 quick wrap-up so you can see the end-to-end flow:

Diagnosis you got right: the cleaned-up subtracted_and_rotated_from + the AdaptImages.updated_via_instance_from rotation of the cached mesh-grid is exactly the fix. Your dev_Q had it correct; the only thing missing was a re-test after the "saving before updates" commit.

Why source reconstruction was weird: on main, the cached galaxy_name_image_plane_mesh_grid_dict was never being shifted by grid_offset, let alone rotated. For adaptive pixelizations in multi-band fits, the cached mesh stayed in the original image-plane frame while FitDataset.grids correctly transformed the data grid — the two ended up in different coordinate frames, the inversion mapped image pixels onto a misaligned source mesh, and the reconstruction looked wrong. Your patch to adapt_images.py resolves it.

Forward-port: all five PRs are now merged on main:

• PyAutoArray#312 — DatasetModel.grid_rotation_angle + Grid2D / Grid2DIrregular subtracted_and_rotated_from.
• PyAutoGalaxy#416 — AdaptImages.updated_via_instance_from rotation of cached mesh grid (your fix).
• PyAutoLens#512 — AnalysisImaging.fit_from plumbing of dataset_model + xp.
• autolens_workspace#154 — multi/features/dataset_offsets/ example extended to inject and recover a 1.5° roll-angle.
• autolens_workspace_test#97 — two parity regression scripts (scripts/multi/dataset_model_parity_lp_linear.py and _delaunay.py) that simulate two datasets of the same scene in different image-plane frames and assert per-dataset log-likelihoods agree between profile-baked and DatasetModel routes. The Delaunay variant agrees to 6 decimal places — that's the long-term guard against this exact bug coming back.

What got dropped: the *MultiBand mass-profile classes (PowerLawMultiBand, PowerLawMultipoleMultiBand, ExternalShearMultiBand, NFWMCRScatterLudlowSphMultiBand) and mge_model_V2_from. With grid-level rotation working, the per-profile workaround isn't needed — Hannah can use any existing mass profile and add grid_rotation_angle as a free parameter on the DatasetModel.

For Hannah's JWST workflow, the pattern is exactly the existing grid_offset one: model_analysis.dataset_model.grid_rotation_angle = af.UniformPrior(lower_limit=-5.0, upper_limit=5.0) on every band after the first. Worked example: autolens_workspace/scripts/multi/features/dataset_offsets/modeling.py.

Total time from Slack message to all-merged: one CLI conversation. Closes this issue.

[Jammy2211]

All five PRs merged. Feature branches deleted, canonical checkouts on main. Closing.