Merge pull request #384 from Jammy2211/feature/fft_jax_imaging

Feature/fft jax imaging

Author: Jammy2211

autolens/__init__.py

@@ -1,10 +1,7 @@
 from autoconf import jax_wrapper
 from autoconf.dictable import from_dict, from_json, output_to_json, to_dict
 from autoarray import preprocess
-from autoarray.dataset.interferometer.w_tilde import (
-    load_curvature_preload_if_compatible,
-)
-from autoarray.dataset.imaging.w_tilde import WTildeImaging
+
 from autoarray.dataset.imaging.dataset import Imaging
 from autoarray.dataset.interferometer.dataset import (
     Interferometer,

autolens/imaging/fit_imaging.py

@@ -112,7 +112,7 @@ def tracer_to_inversion(self) -> TracerToInversion:
             noise_map=self.noise_map,
             grids=self.grids,
             psf=self.dataset.psf,
-            w_tilde=self.w_tilde,
+            sparse_operator=self.dataset.sparse_operator,
         )
 
         return TracerToInversion(

autolens/interferometer/fit_interferometer.py

@@ -115,7 +115,7 @@ def tracer_to_inversion(self) -> TracerToInversion:
             noise_map=self.noise_map,
             grids=self.grids,
             transformer=self.dataset.transformer,
-            w_tilde=self.w_tilde,
+            sparse_operator=self.dataset.sparse_operator,
         )
 
         return TracerToInversion(

autolens/lens/to_inversion.py

@@ -175,18 +175,13 @@ def lp_linear_func_list_galaxy_dict(
                 blurring=traced_blurring_grids_of_planes_list[plane_index],
             )
 
-            if self.dataset.w_tilde is not None:
-                w_tilde = self.dataset.w_tilde
-            else:
-                w_tilde = None
-
             dataset = aa.DatasetInterface(
                 data=self.dataset.data,
                 noise_map=self.dataset.noise_map,
                 grids=grids,
                 psf=self.psf,
                 transformer=self.transformer,
-                w_tilde=w_tilde,
+                sparse_operator=self.dataset.sparse_operator,
             )
 
             galaxies_to_inversion = ag.GalaxiesToInversion(

autolens/lens/tracer_util.py

@@ -172,7 +172,9 @@ def traced_grid_2d_list_from(
         )
 
         # Remove NaN deflection values to sanitize the ray-tracing calculation for JAX.
-        deflections_yx_2d = xp.where(xp.isfinite(deflections_yx_2d.array), deflections_yx_2d.array, 0.0)
+        deflections_yx_2d = xp.where(
+            xp.isfinite(deflections_yx_2d.array), deflections_yx_2d.array, 0.0
+        )
 
         traced_deflection_list.append(deflections_yx_2d)
 
@@ -348,9 +350,7 @@ def time_delays_from(
     )
 
     # Time-delay distance in meters: (1+z_l) * Dd * Ds / Dds
-    D_dt_m = (
-        (1.0 + z_l) * (Dd_kpc * Ds_kpc / Dds_kpc) * kpc_in_m
-    )
+    D_dt_m = (1.0 + z_l) * (Dd_kpc * Ds_kpc / Dds_kpc) * kpc_in_m
 
     # Fermat potential (should be in arcsec^2 for this formula)
     fermat_potential = galaxies.fermat_potential_from(grid=grid, xp=xp)

test_autolens/imaging/test_simulate_and_fit_imaging.py

@@ -651,7 +651,7 @@ def test__simulate_imaging_data_and_fit__linear_light_profiles_and_pixelization_
     assert fit_linear.figure_of_merit == pytest.approx(-180.8284970580511, 1.0e-4)
 
 
-def test__simulate_imaging_data_and_fit__complex_fit_compare_mapping_matrix_w_tilde():
+def test__simulate_imaging_data_and_fit__complex_fit_compare_mapping_matrix_sparse_operator():
 
     grid = al.Grid2D.uniform(shape_native=(21, 21), pixel_scales=0.1)
 
@@ -735,20 +735,20 @@ def test__simulate_imaging_data_and_fit__complex_fit_compare_mapping_matrix_w_ti
         tracer=tracer,
     )
 
-    masked_dataset_w_tilde = masked_dataset.apply_w_tilde()
+    masked_dataset_sparse_operator = masked_dataset.apply_sparse_operator()
 
-    fit_w_tilde = al.FitImaging(
-        dataset=masked_dataset_w_tilde,
+    fit_sparse_operator = al.FitImaging(
+        dataset=masked_dataset_sparse_operator,
         tracer=tracer,
     )
 
     assert fit_mapping.inversion.curvature_matrix == pytest.approx(
-            fit_w_tilde.inversion.curvature_matrix,
+            fit_sparse_operator.inversion.curvature_matrix,
         1.0e-4,
     )
 
     assert fit_mapping.inversion.regularization_matrix == pytest.approx(
-            fit_w_tilde.inversion.regularization_matrix,
+            fit_sparse_operator.inversion.regularization_matrix,
         1.0e-4,
     )