Merge pull request #387 from Jammy2211/feature/jax_in_image_dict

Feature/jax in image dict

Author: Jammy2211

autolens/imaging/fit_imaging.py

@@ -172,6 +172,7 @@ def galaxy_image_dict(self) -> Dict[ag.Galaxy, np.ndarray]:
 
         galaxy_image_2d_dict = self.tracer.galaxy_image_2d_dict_from(
             grid=self.grids.lp,
+            xp=self._xp
         )
 
         galaxy_linear_obj_image_dict = self.galaxy_linear_obj_data_dict_from(
@@ -199,6 +200,7 @@ def galaxy_model_image_dict(self) -> Dict[ag.Galaxy, np.ndarray]:
             grid=self.grids.lp,
             psf=self.dataset.psf,
             blurring_grid=self.grids.blurring,
+            xp=self._xp
         )
 
         galaxy_linear_obj_image_dict = self.galaxy_linear_obj_data_dict_from(

autolens/interferometer/fit_interferometer.py

@@ -173,7 +173,9 @@ def galaxy_image_dict(self) -> Dict[ag.Galaxy, np.ndarray]:
         For modeling, this dictionary is used to set up the `adapt_images` that adapt certain pixelizations to the
         data being fitted.
         """
-        galaxy_image_dict = self.tracer.galaxy_image_2d_dict_from(grid=self.grids.lp)
+        galaxy_image_dict = self.tracer.galaxy_image_2d_dict_from(
+            grid=self.grids.lp, xp=self._xp
+        )
 
         galaxy_linear_obj_image_dict = self.galaxy_linear_obj_data_dict_from(
             use_operated=False
@@ -194,7 +196,7 @@ def galaxy_model_visibilities_dict(self) -> Dict[ag.Galaxy, np.ndarray]:
           are solved for first via the inversion.
         """
         galaxy_model_visibilities_dict = self.tracer.galaxy_visibilities_dict_from(
-            grid=self.grids.lp, transformer=self.dataset.transformer
+            grid=self.grids.lp, transformer=self.dataset.transformer, xp=self._xp
         )
 
         galaxy_linear_obj_visibilities_dict = self.galaxy_linear_obj_data_dict_from(

autolens/lens/tracer.py

@@ -579,7 +579,7 @@ def image_2d_via_input_plane_image_from(
         return aa.Array2D(values=image, mask=grid.mask, xp=xp)
 
     def galaxy_image_2d_dict_from(
-        self, grid: aa.type.Grid2DLike, operated_only: Optional[bool] = None
+        self, grid: aa.type.Grid2DLike, xp=np, operated_only: Optional[bool] = None
     ) -> Dict[ag.Galaxy, np.ndarray]:
         """
         Returns a dictionary associating every `Galaxy` object in the `Tracer` with its corresponding 2D image, using
@@ -605,12 +605,14 @@ def galaxy_image_2d_dict_from(
 
         galaxy_image_2d_dict = dict()
 
-        traced_grid_list = self.traced_grid_2d_list_from(grid=grid)
+        traced_grid_list = self.traced_grid_2d_list_from(grid=grid, xp=xp)
 
         for plane_index, galaxies in enumerate(self.planes):
             image_2d_list = [
                 galaxy.image_2d_from(
-                    grid=traced_grid_list[plane_index], operated_only=operated_only
+                    grid=traced_grid_list[plane_index],
+                    operated_only=operated_only,
+                    xp=xp,
                 )
                 for galaxy in galaxies
             ]

test_autolens/interferometer/test_fit_interferometer.py

@@ -180,12 +180,8 @@ def test___galaxy_image_dict(interferometer_7, interferometer_7_grid):
         tracer=tracer,
     )
 
-    assert fit.galaxy_image_dict[g0_linear][4] == pytest.approx(
-        1.00018622848, 1.0e-2
-    )
-    assert fit.galaxy_image_dict[g1_linear][3] == pytest.approx(
-        -0.017435532289, 1.0e-2
-    )
+    assert fit.galaxy_image_dict[g0_linear][4] == pytest.approx(1.00018622848, 1.0e-2)
+    assert fit.galaxy_image_dict[g1_linear][3] == pytest.approx(-0.017435532289, 1.0e-2)
 
     pixelization = al.Pixelization(
         mesh=al.mesh.RectangularUniform(shape=(3, 3)),