fix: Basis.image_2d_from and dPIEPotential.convergence_2d_from return wrong wrapper types

[Jammy2211]

Overview

Two library bugs surfaced while writing the new scripts/guides/profiles/{light,mass}.py guides in autolens_workspace (#86 / #176 / #178 / #179):

1. Basis.image_2d_from returns a raw numpy.ndarray when the basis is composed of LightProfileLinear constituents (the standard MGE case), instead of the Array2D its return-type annotation promises.
2. dPIEPotential.convergence_2d_from returns a VectorYX2D instead of the Array2D its scalar return value should produce.

Both forced workarounds in the workspace guides — neither is structurally serious, but both are tiny library-side fixes that remove the need for those workarounds and tighten the type contract advertised by the methods' annotations.

Plan

• Fix 1 — Basis.image_2d_from: in autogalaxy/profiles/basis.py, wrap the LightProfileLinear zero placeholder in aa.Array2D so the summed return type is always Array2D, regardless of which constituent profiles the basis was built from.
• Fix 2 — dPIEPotential.convergence_2d_from: in autogalaxy/profiles/mass/total/dual_pseudo_isothermal_potential.py, swap the @aa.decorators.to_vector_yx decorator on dPIEPotential.convergence_2d_from for @aa.decorators.to_array. The function body returns a scalar field; the vector decorator is a copy-paste typo from the deflections method directly above it.
• Regression tests: add one test per fix asserting the return type is Array2D for each method, in the existing test_autogalaxy/profiles/test_basis.py and test_autogalaxy/profiles/mass/total/test_dual_pseudo_isothermal_potential.py.

Detailed implementation plan

Affected Repositories

• PyAutoGalaxy (primary)
• PyAutoLens — none (re-exports only; will pick up the fix transparently)

Work Classification

Library (code + tests). Ships via /ship_library.

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoGalaxy	main	clean
./PyAutoLens	main	clean (reference; not modified)

Suggested branch: feature/profile-return-type-fixes
Worktree root: ~/Code/PyAutoLabs-wt/profile-return-type-fixes/
Routing: library-dev → /start_library.

Implementation Steps

1. autogalaxy/profiles/basis.py — In Basis.image_2d_list_from, replace the LightProfileLinear placeholder

   xp.zeros((grid.shape[0],))

   with

   aa.Array2D(values=xp.zeros((grid.shape[0],)), mask=grid.mask)

   This makes Basis.image_2d_from = sum([...]) return an Array2D whenever every constituent is linear (the MGE case), because Array2D.__add__ returns Array2D. The standard-profile path already produces Array2D instances from each image_2d_from call, so the sum is unchanged for non-linear bases.

2. autogalaxy/profiles/mass/total/dual_pseudo_isothermal_potential.py — Change the decorator on dPIEPotential.convergence_2d_from from:

   @aa.decorators.to_vector_yx
   @aa.decorators.transform
   def convergence_2d_from(...):

   to:

   @aa.decorators.to_array
   @aa.decorators.transform
   def convergence_2d_from(...):

   Matches the dPIEPotentialSph.convergence_2d_from decorator stack directly below it, which has always been correct.

3. Add regression tests in the existing test files:

   • test_autogalaxy/profiles/test_basis.py — add a test that constructs a Basis of LightProfileLinear Gaussians and asserts isinstance(basis.image_2d_from(grid), aa.Array2D).
   • test_autogalaxy/profiles/mass/total/test_dual_pseudo_isothermal_potential.py — add a test that constructs a dPIEPotential and asserts isinstance(profile.convergence_2d_from(grid), aa.Array2D). (NOT VectorYX2D.)

4. Run the existing test suite to confirm no regressions on the surrounding behaviour:

   python -m pytest test_autogalaxy/profiles/test_basis.py test_autogalaxy/profiles/mass/total/test_dual_pseudo_isothermal_potential.py -x

Key Files

• autogalaxy/profiles/basis.py — wrap linear-profile zero placeholder in Array2D
• autogalaxy/profiles/mass/total/dual_pseudo_isothermal_potential.py — fix decorator
• test_autogalaxy/profiles/test_basis.py — add regression test
• test_autogalaxy/profiles/mass/total/test_dual_pseudo_isothermal_potential.py — add regression test

Out of scope

• Workspace cleanup of the workarounds in autolens_workspace/scripts/guides/profiles/{light,mass}.py — cosmetic, can stay until the next pass through those scripts.
• Any other lp_basis / dPIE-family quirks not specifically listed here.

JAX considerations

• Basis.image_2d_from is consumed by downstream Python (Galaxy → Galaxies → Tracer), never the JIT output, so wrapping in Array2D is safe even in xp=jnp mode (per the CLAUDE.md rule on intermediate-step Array2D wraps).
• dPIEPotential.convergence_2d_from is decorated; the decorator stack handles the JAX-vs-numpy path, no extra work needed.

Original Prompt

Click to expand starting prompt

Two library bugs surfaced while writing autolens_workspace/scripts/guides/profiles/light.py (#86 / #176) and mass.py (#178 / #179):

1. Basis.image_2d_from returns a raw numpy.ndarray instead of an Array2D when the basis is composed of LightProfileLinear constituents. The light.py guide had to wrap the basis in a Galaxy and plot via the galaxy's image to make aplt.plot_array work.

   Root cause: Basis.image_2d_list_from returns xp.zeros((grid.shape[0],)) for LightProfileLinear profiles (because their intensity will be solved by the inversion later). Summing those zero ndarrays in image_2d_from produces a raw ndarray rather than an Array2D.

2. dPIEPotential.convergence_2d_from returns a VectorYX2D instead of an Array2D. The mass.py guide had to skip plotting convergence for dPIEPotential and use dPIEPotentialSph in the walkthrough instead.

   Root cause: a typo / copy-paste of the deflections decorator — @aa.decorators.to_vector_yx is used where @aa.decorators.to_array is correct. The function body returns a scalar field (kappa_circ * (1 - asymm_term) + (alpha_circ / grid_radii) * asymm_term) which is not a (y, x) vector.

Both fixes are tiny library-code changes with regression tests.

[Jammy2211]

Shipped — 2026-05-18

PR

• fix: Basis.image_2d_from and dPIEPotential.convergence_2d_from return wrong wrapper types #425 (merged)

Summary

Fixed two profile-return-type bugs that surfaced while writing the autolens_workspace mass / light profile guides (#86 / #176 / #178 / #179):

1. Basis.image_2d_from now returns aa.Array2D uniformly, even when the basis is composed entirely of LightProfileLinear constituents (the MGE case). Previously the linear-profile zero placeholder was a raw xp.zeros((N,)) ndarray, making the summed image_2d_from result a raw ndarray.
2. dPIEPotential.convergence_2d_from now returns aa.Array2D instead of aa.VectorYX2D. The decorator was a copy-paste typo from the deflections method directly above — swapped @aa.decorators.to_vector_yx for @aa.decorators.to_array.

Regression tests added for both fixes. Full 911-test test_autogalaxy/ suite passes.

Session Notes

• No workspace migration needed: the workarounds in light.py and mass.py (Galaxy-wrap the Basis, plot dPIEPotentialSph instead of dPIEPotential in the walkthrough) are now technically obsolete but stay in place as cosmetic safety. They can come out on the next pass through those scripts.
• No API surface change: both methods already advertised aa.Array2D return types in their annotations; this PR only makes runtime behaviour match. Any external caller depending on the (incorrect) ndarray / VectorYX2D returns will see a tighter type — none in the codebase or workspaces do, since both methods were effectively unusable through aplt.plot_array before this fix.