Feature/jax unit tests

autogalaxy/config/notation.yaml

@@ -72,7 +72,6 @@ label:
     SMBH: smbh
     Redshift: z
     Regularization: reg
-    InputDeflections: defl
 label_format:
   format:
     sigma: '{:.4f}'

autogalaxy/galaxy/galaxy.py

@@ -1,10 +1,12 @@
 from typing import Dict, List, Optional, Type, Union
 
+import jax.numpy as jnp
 import numpy as np
 
+from autoconf.dictable import instance_as_dict, to_dict
+
 import autoarray as aa
 import autofit as af
-from autoconf.dictable import instance_as_dict, to_dict
 
 from autogalaxy import exc
 from autogalaxy.operate.deflections import OperateDeflections
@@ -228,7 +230,7 @@ def image_2d_from(
         ):
             return sum(self.image_2d_list_from(grid=grid, operated_only=operated_only))
 
-        return np.zeros((grid.shape[0],))
+        return jnp.zeros((grid.shape[0],))
 
     @aa.grid_dec.to_projected
     def image_1d_from(self, grid: aa.type.Grid2DLike) -> np.ndarray:
@@ -261,7 +263,7 @@ def image_1d_from(self, grid: aa.type.Grid2DLike) -> np.ndarray:
 
             return sum(image_1d_list)
 
-        return np.zeros((grid.shape[0],))
+        return jnp.zeros((grid.shape[0],))
 
     @aa.grid_dec.to_vector_yx
     def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarray:
@@ -289,7 +291,7 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarr
                     self.cls_list_from(cls=MassProfile),
                 )
             )
-        return np.zeros((grid.shape[0], 2))
+        return jnp.zeros((grid.shape[0], 2))
 
     @aa.grid_dec.to_array
     def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarray:
@@ -316,7 +318,7 @@ def convergence_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarray:
                     self.cls_list_from(cls=MassProfile),
                 )
             )
-        return np.zeros((grid.shape[0],))
+        return jnp.zeros((grid.shape[0],))
 
     @aa.grid_dec.to_grid
     def traced_grid_2d_from(self, grid: aa.type.Grid2DLike) -> aa.type.Grid2DLike:
@@ -355,6 +357,7 @@ def convergence_1d_from(self, grid: aa.type.Grid1D2DLike) -> np.ndarray:
             convergence_1d_list = []
 
             for mass_profile in self.cls_list_from(cls=MassProfile):
+
                 grid_radial = self.grid_radial_from(
                     grid=grid, centre=mass_profile.centre, angle=mass_profile.angle
                 )
@@ -365,7 +368,7 @@ def convergence_1d_from(self, grid: aa.type.Grid1D2DLike) -> np.ndarray:
 
             return sum(convergence_1d_list)
 
-        return np.zeros((grid.shape[0],))
+        return jnp.zeros((grid.shape[0],))
 
     @aa.grid_dec.to_array
     def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarray:
@@ -392,7 +395,7 @@ def potential_2d_from(self, grid: aa.type.Grid2DLike, **kwargs) -> np.ndarray:
                     self.cls_list_from(cls=MassProfile),
                 )
             )
-        return np.zeros((grid.shape[0],))
+        return jnp.zeros((grid.shape[0],))
 
     @aa.grid_dec.to_projected
     def potential_1d_from(self, grid: aa.type.Grid2DLike) -> np.ndarray:
@@ -425,7 +428,7 @@ def potential_1d_from(self, grid: aa.type.Grid2DLike) -> np.ndarray:
 
             return sum(potential_1d_list)
 
-        return np.zeros((grid.shape[0],))
+        return jnp.zeros((grid.shape[0],))
 
     @property
     def half_light_radius(self):

autogalaxy/imaging/fit_imaging.py

@@ -121,6 +121,7 @@ def galaxies_to_inversion(self) -> GalaxiesToInversion:
             noise_map=self.noise_map,
             grids=self.grids,
             psf=self.dataset.psf,
+            convolver=self.dataset.convolver,
             w_tilde=self.w_tilde,
         )