Revert "Feature/jax point source"

Author: Jammy2211

autolens/__init__.py

@@ -10,6 +10,7 @@
 from autoarray.mask.mask_1d import Mask1D
 from autoarray.mask.mask_2d import Mask2D
 from autoarray.mask.derive.zoom_2d import Zoom2D
+from autoarray.operators.convolver import Convolver
 from autoarray.operators.over_sampling.over_sampler import OverSampler  # noqa
 from autoarray.inversion.inversion.dataset_interface import DatasetInterface
 from autoarray.inversion.inversion.mapper_valued import MapperValued
@@ -26,7 +27,6 @@
 from autoarray.inversion.pixelization.mappers.mapper_grids import MapperGrids
 from autoarray.inversion.pixelization.mappers.factory import mapper_from as Mapper
 from autoarray.inversion.pixelization.border_relocator import BorderRelocator
-from autoarray.inversion.convolver import Convolver
 from autoarray.operators.transformer import TransformerDFT
 from autoarray.operators.transformer import TransformerNUFFT
 from autoarray.structures.arrays.uniform_1d import Array1D

autolens/analysis/positions.py

@@ -71,7 +71,6 @@ def __init__(
             The plane redshift of the lensed source multiple images, which is only required if position threshold
             for a double source plane lens system is being used where the specific plane is required.
         """
-
         self.positions = positions
         self.threshold = threshold
         self.plane_redshift = plane_redshift
@@ -166,12 +165,10 @@ def log_likelihood_penalty_base_from(
                 residual_map=residual_map, noise_map=dataset.noise_map
             )
 
-            chi_squared = aa.util.fit.chi_squared_from(
-                chi_squared_map=chi_squared_map.array
-            )
+            chi_squared = aa.util.fit.chi_squared_from(chi_squared_map=chi_squared_map)
 
             noise_normalization = aa.util.fit.noise_normalization_from(
-                noise_map=dataset.noise_map.array
+                noise_map=dataset.noise_map
             )
 
         else:

autolens/analysis/result.py

@@ -233,7 +233,7 @@ def positions_threshold_from(
             data=positions, noise_map=None, tracer=tracer, plane_redshift=plane_redshift
         )
 
-        threshold = factor * np.nanmax(positions_fits.max_separation_of_plane_positions)
+        threshold = factor * np.max(positions_fits.max_separation_of_plane_positions)
 
         if minimum_threshold is not None:
             if threshold < minimum_threshold:
@@ -285,7 +285,7 @@ def positions_likelihood_from(
 
         Returns
         -------
-        The `PositionsLH` object used to apply a likelihood penalty or resample the positions.
+        The `PositionsLH` object used to apply a likelihood penalty using the positions.
         """
 
         if os.environ.get("PYAUTOFIT_TEST_MODE") == "1":
@@ -308,12 +308,6 @@ def positions_likelihood_from(
 
             positions = positions[distances > mass_centre_radial_distance_min]
 
-        mask = np.isfinite(positions.array).all(axis=1)
-
-        positions = aa.Grid2DIrregular(
-            positions[mask]
-        )
-
         threshold = self.positions_threshold_from(
             factor=factor,
             minimum_threshold=minimum_threshold,

autolens/imaging/fit_imaging.py

@@ -82,17 +82,9 @@ 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,
+            convolver=self.dataset.convolver,
             blurring_grid=self.grids.blurring,
         )
 
@@ -101,6 +93,7 @@ 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
@@ -110,7 +103,6 @@ def tracer_to_inversion(self) -> TracerToInversion:
             data=self.profile_subtracted_image,
             noise_map=self.noise_map,
             grids=self.grids,
-            psf=self.dataset.psf,
             convolver=self.dataset.convolver,
             w_tilde=self.w_tilde,
         )
@@ -171,7 +163,7 @@ def galaxy_model_image_dict(self) -> Dict[ag.Galaxy, np.ndarray]:
 
         galaxy_blurred_image_2d_dict = self.tracer.galaxy_blurred_image_2d_dict_from(
             grid=self.grids.lp,
-            psf=self.dataset.psf,
+            convolver=self.dataset.convolver,
             blurring_grid=self.grids.blurring,
         )
 

autolens/imaging/model/analysis.py

@@ -117,8 +117,6 @@ 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].array)
+                self.mat_plot_2d.cmap.kwargs["vmax"] = np.max(self.fit.model_images_of_planes_list[plane_index])
 
             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([model_image.array for model_image in self.fit.model_images_of_planes_list[1:]])
+                self.mat_plot_2d.cmap.kwargs["vmax"] = np.max(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([model_image.array for model_image in self.fit.model_images_of_planes_list[1:]])
+                self.mat_plot_2d.cmap.kwargs["vmax"] = np.max(self.fit.model_images_of_planes_list[1:])
 
             self.mat_plot_2d.plot_array(
                 array=self.fit.model_data,

autolens/lens/to_inversion.py

@@ -180,8 +180,7 @@ def lp_linear_func_list_galaxy_dict(
                 data=self.dataset.data,
                 noise_map=self.dataset.noise_map,
                 grids=grids,
-                psf=self.psf,
-                convolver=self.dataset.convolver,
+                convolver=self.convolver,
                 transformer=self.transformer,
                 w_tilde=self.dataset.w_tilde,
             )

autolens/lens/tracer.py

@@ -1,5 +1,6 @@
 from abc import ABC
 import numpy as np
+from functools import wraps
 from scipy.interpolate import griddata
 from typing import Dict, List, Optional, Type, Union
 
@@ -548,9 +549,9 @@ def image_2d_via_input_plane_image_from(
         )[plane_index]
 
         image = griddata(
-            points=plane_grid.array,
-            values=plane_image.array,
-            xi=traced_grid.over_sampled.array,
+            points=plane_grid,
+            values=plane_image,
+            xi=traced_grid.over_sampled,
             fill_value=0.0,
             method="linear",
         )
@@ -1190,5 +1191,5 @@ def set_snr_of_snr_light_profiles(
                         )
 
     @aa.profile_func
-    def convolve_via_psf(self, image, blurring_image, psf):
-        return psf.convolve_image(image=image, blurring_image=blurring_image)
+    def convolve_via_convolver(self, image, blurring_image, convolver):
+        return convolver.convolve_image(image=image, blurring_image=blurring_image)

autolens/point/fit/fluxes.py

@@ -1,4 +1,3 @@
-import jax.numpy as jnp
 from typing import Optional
 
 import autoarray as aa
@@ -102,10 +101,10 @@ def model_data(self):
         are used.
         """
         return aa.ArrayIrregular(
-            values=jnp.array([
+            values=[
                 magnification * self.profile.flux
                 for magnification in self.magnifications_at_positions
-            ])
+            ]
         )
 
     @property
@@ -129,5 +128,5 @@ def chi_squared(self) -> float:
         RMS noise-map values squared.
         """
         return ag.util.fit.chi_squared_from(
-            chi_squared_map=self.chi_squared_map.array,
+            chi_squared_map=self.chi_squared_map,
         )

autolens/point/fit/positions/image/pair.py

@@ -1,7 +1,4 @@
-import jax.numpy as jnp
 import numpy as np
-from ott.geometry import pointcloud
-from ott.solvers.linear import sinkhorn
 from scipy.optimize import linear_sum_assignment
 
 import autoarray as aa