Feature/jax linear light

autolens/imaging/fit_imaging.py

@@ -82,6 +82,14 @@ def blurred_image(self) -> aa.Array2D:
         """
         Returns the image of all light profiles in the fit's tracer convolved with the imaging dataset's PSF.
         """
+
+        if len(self.tracer.cls_list_from(cls=ag.LightProfile)) == len(
+            self.tracer.cls_list_from(cls=ag.lp_operated.LightProfileOperated)
+        ):
+            return self.tracer.image_2d_from(
+                grid=self.grids.lp,
+            )
+
         return self.tracer.blurred_image_2d_from(
             grid=self.grids.lp,
             psf=self.dataset.psf,
@@ -93,7 +101,6 @@ def profile_subtracted_image(self) -> aa.Array2D:
         """
         Returns the dataset's image with all blurred light profile images in the fit's tracer subtracted.
         """
-
         return self.data - self.blurred_image
 
     @property

autolens/imaging/model/analysis.py

@@ -117,6 +117,8 @@ def log_likelihood_function(self, instance: af.ModelInstance) -> float:
             np.linalg.LinAlgError,
             OverflowError,
         ) as e:
+            print(e)
+            fggdfg
             raise exc.FitException from e
 
     def fit_from(

autolens/imaging/plot/fit_imaging_plotters.py

@@ -216,7 +216,7 @@ def figures_2d_of_planes(
         for plane_index in plane_indexes:
 
             if use_source_vmax:
-                self.mat_plot_2d.cmap.kwargs["vmax"] = np.max(self.fit.model_images_of_planes_list[plane_index])
+                self.mat_plot_2d.cmap.kwargs["vmax"] = np.max(self.fit.model_images_of_planes_list[plane_index].array)
 
             if subtracted_image:
 
@@ -765,7 +765,7 @@ def figures_2d(
         if data:
 
             if use_source_vmax:
-                self.mat_plot_2d.cmap.kwargs["vmax"] = np.max(self.fit.model_images_of_planes_list[1:])
+                self.mat_plot_2d.cmap.kwargs["vmax"] = np.max([model_image.array for model_image in self.fit.model_images_of_planes_list[1:]])
 
             self.mat_plot_2d.plot_array(
                 array=self.fit.data,
@@ -799,7 +799,7 @@ def figures_2d(
         if model_image:
 
             if use_source_vmax:
-                self.mat_plot_2d.cmap.kwargs["vmax"] = np.max(self.fit.model_images_of_planes_list[1:])
+                self.mat_plot_2d.cmap.kwargs["vmax"] = np.max([model_image.array for model_image in self.fit.model_images_of_planes_list[1:]])
 
             self.mat_plot_2d.plot_array(
                 array=self.fit.model_data,

autolens/point/fit/times_delays.py

@@ -102,5 +102,5 @@ def chi_squared(self) -> float:
         which is the residual values divided by the RMS noise-map squared.
         """
         return ag.util.fit.chi_squared_from(
-            chi_squared_map=self.chi_squared_map,
+            chi_squared_map=self.chi_squared_map.array,
         )

autolens/point/max_separation.py

@@ -42,7 +42,7 @@ def __init__(
         except TypeError:
             plane_index = -1
 
-        self.plane_positions = tracer.traced_grid_2d_list_from(grid=data)[plane_index]
+        self.plane_positions = aa.Grid2DIrregular(values=tracer.traced_grid_2d_list_from(grid=data)[plane_index])
 
     @property
     def furthest_separations_of_plane_positions(self) -> aa.ArrayIrregular:

test_autolens/imaging/test_fit_imaging.py

@@ -519,7 +519,7 @@ def test__galaxy_model_image_dict(masked_imaging_7x7):
     assert (fit.galaxy_model_image_dict[g2] == np.zeros(9)).all()
 
     assert fit.model_data == pytest.approx(
-        fit.galaxy_model_image_dict[g0_linear] + fit.galaxy_model_image_dict[g1_linear],
+        fit.galaxy_model_image_dict[g0_linear].array + fit.galaxy_model_image_dict[g1_linear].array,
         1.0e-4,
     )
 
@@ -540,13 +540,13 @@ def test__galaxy_model_image_dict(masked_imaging_7x7):
 
     fit = al.FitImaging(dataset=masked_imaging_7x7, tracer=tracer)
 
-    assert (fit.galaxy_model_image_dict[g0_no_light] == np.zeros(9)).all()
-    assert fit.galaxy_model_image_dict[galaxy_pix_0][4] == pytest.approx(
+    assert (fit.galaxy_model_image_dict[g0_no_light].array == np.zeros(9)).all()
+    assert fit.galaxy_model_image_dict[galaxy_pix_0].array[4] == pytest.approx(
         1.259965886, 1.0e-4
     )
 
     assert fit.model_data == pytest.approx(
-        fit.galaxy_model_image_dict[galaxy_pix_0], 1.0e-4
+        fit.galaxy_model_image_dict[galaxy_pix_0].array, 1.0e-4
     )
 
     # Normal light + Linear Light PRofiles + Pixelization + Regularizaiton

test_autolens/imaging/test_simulate_and_fit_imaging.py

@@ -360,7 +360,7 @@ def test__simulate_imaging_data_and_fit__linear_light_profiles_and_pixelization(
                 0.0,
             ]
         ),
-        1.0e-4,
+        abs=1.0e-1,
     )
     assert fit_linear.figure_of_merit == pytest.approx(-84.11166, 1.0e-4)
 
@@ -382,7 +382,7 @@ def test__simulate_imaging_data_and_fit__linear_light_profiles_and_pixelization(
         settings_inversion=al.SettingsInversion(use_w_tilde=False),
     )
 
-    assert fit_linear.figure_of_merit == pytest.approx(-73.27676850869975, 1.0e-4)
+    assert fit_linear.figure_of_merit == pytest.approx(-73.27676850869975, abs=1.0e-4)
 
 
 def test__simulate_imaging_data_and_fit__linear_light_profiles_and_pixelization__sub_2():
@@ -512,9 +512,9 @@ def test__simulate_imaging_data_and_fit__linear_light_profiles_and_pixelization_
                 0.0,
             ]
         ),
-        1.0e-4,
+        abs=1.0e-1,
     )
-    assert fit_linear.figure_of_merit == pytest.approx(-84.66302233089499, 1.0e-4)
+    assert fit_linear.figure_of_merit == pytest.approx(-84.66302233089499, abs=1.0e-4)
 
 
 def test__simulate_imaging_data_and_fit__complex_fit_compare_mapping_matrix_w_tilde():