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

autogalaxy/profiles/basis.py

@@ -147,7 +147,9 @@ def image_2d_list_from(
                     grid=grid, xp=xp, operated_only=operated_only
                 )
                 if not isinstance(light_profile, lp_linear.LightProfileLinear)
-                else xp.zeros((grid.shape[0],))
+                else aa.Array2D(
+                    values=xp.zeros((grid.shape[0],)), mask=grid.mask
+                )
             )
             for light_profile in self.light_profile_list
         ]

autogalaxy/profiles/mass/total/dual_pseudo_isothermal_potential.py

@@ -118,7 +118,7 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
 
         return xp.vstack((deflection_y, deflection_x)).T
 
-    @aa.decorators.to_vector_yx
+    @aa.decorators.to_array
     @aa.decorators.transform
     def convergence_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs):
         """

test_autogalaxy/profiles/mass/total/test_dual_pseudo_isothermal_potential.py

@@ -1,10 +1,23 @@
 import pytest
 
+import autoarray as aa
 import autogalaxy as ag
 
 grid = ag.Grid2DIrregular([[1.0, 1.0], [2.0, 2.0], [3.0, 3.0], [2.0, 4.0]])
 
 
+def test__convergence_2d_from__returns_array2d_not_vector():
+    # Regression test for a copy-paste typo: convergence_2d_from used to be
+    # decorated with @aa.decorators.to_vector_yx (the deflections decorator
+    # directly above it) which wrapped the scalar convergence in a VectorYX2D.
+    # The correct decorator is @aa.decorators.to_array.
+    mp = ag.mp.dPIEPotential(centre=(0.0, 0.0), ell_comps=(0.05, 0.0), ra=0.2, rs=2.0, b0=1.0)
+
+    convergence = mp.convergence_2d_from(grid=grid)
+
+    assert isinstance(convergence, aa.ArrayIrregular)
+
+
 def test__deflections_yx_2d_from__sph_config_1():
     mp = ag.mp.dPIEPotentialSph(centre=(-0.7, 0.5), b0=5.2, ra=2.0, rs=3.0)
 

test_autogalaxy/profiles/test_basis.py

@@ -1,6 +1,7 @@
 import numpy as np
 import pytest
 
+import autoarray as aa
 import autogalaxy as ag
 
 
@@ -42,6 +43,24 @@ def test__image_2d_from__does_not_include_linear_light_profiles(grid_2d_7x7):
     assert (image == lp_image).all()
 
 
+def test__image_2d_from__returns_array2d_for_linear_only_basis(grid_2d_7x7):
+    # Regression test for the case where every constituent is a LightProfileLinear:
+    # the placeholder for the linear-intensity-unknown image used to be a raw
+    # ndarray of zeros, which made `sum(...)` return a raw ndarray rather than an
+    # `Array2D`.  After the fix the placeholder is itself an `Array2D` so the
+    # summed return type is uniform regardless of the constituent mix.
+    basis = ag.lp_basis.Basis(
+        profile_list=[
+            ag.lp_linear.Gaussian(sigma=0.5),
+            ag.lp_linear.Gaussian(sigma=1.0),
+        ]
+    )
+
+    image = basis.image_2d_from(grid=grid_2d_7x7)
+
+    assert isinstance(image, aa.Array2D)
+
+
 def test__image_2d_from__operated_only_false__returns_only_non_operated_profile_image(
     grid_2d_7x7, lp_0, lp_operated_0
 ):