Feature/fft jax imaging

autogalaxy/__init__.py

@@ -7,9 +7,6 @@
 from autoconf.dictable import from_dict, from_json, output_to_json, to_dict
 from autoarray.dataset import preprocess  # noqa
 
-# from autoarray.dataset.interferometer.w_tilde import (
-#     load_curvature_preload_if_compatible,
-# )
 from autoarray.dataset.imaging.dataset import Imaging  # noqa
 from autoarray.dataset.interferometer.dataset import Interferometer  # noqa
 from autoarray.dataset.dataset_model import DatasetModel

autogalaxy/abstract_fit.py

@@ -65,18 +65,18 @@ def perform_inversion(self) -> bool:
         return self.model_obj.perform_inversion
 
     @property
-    def w_tilde(self) -> Optional[aa.WTildeImaging]:
+    def sparse_operator(self) -> Optional[aa.ImagingSparseOperator]:
         """
-        Only call the `w_tilde` property of a dataset used to perform efficient linear algebra calcualtions if
-        the SettingsInversion()` object has `use_w_tilde=True`, to avoid unnecessary computation.
+        Only call the `sparse_operator` property of a dataset used to perform efficient linear algebra calculations if
+        the SettingsInversion()` object has `use_sparse_operator=True`, to avoid unnecessary computation.
 
         Returns
         -------
-        The w-tilde matrix if the w-tilde formalism is being used for this inversion.
+        The sparse operator used for efficient linear algebra calculations for this inversion, if enabled.
         """
-        if self.dataset.w_tilde is not None:
+        if self.dataset.sparse_operator is not None:
             if self.total_mappers > 0:
-                return self.dataset.w_tilde
+                return self.dataset.sparse_operator
 
     @property
     def inversion(self) -> Optional[aa.Inversion]:

autogalaxy/aggregator/agg_util.py

@@ -230,9 +230,6 @@ def preloads_from(
     in order to perform the fit. The main purpose is to use the same image-plane mesh centres for pixelization where
     the mesh is computed in the image-plane (see `agg_util.mesh_grids_of_planes_list_from`).
 
-    The preloads may also switch off `w_tilde` so fits are computed faster locally as they do not need to recompute
-    w_tilde.
-
     Parameters
     ----------
     preloads_cls

autogalaxy/imaging/fit_imaging.py

@@ -123,7 +123,7 @@ def galaxies_to_inversion(self) -> GalaxiesToInversion:
             noise_map=self.noise_map,
             grids=self.grids,
             psf=self.dataset.psf,
-            w_tilde=self.w_tilde,
+            sparse_operator=self.dataset.sparse_operator,
         )
 
         return GalaxiesToInversion(

autogalaxy/interferometer/fit_interferometer.py

@@ -103,7 +103,7 @@ def galaxies_to_inversion(self) -> GalaxiesToInversion:
             noise_map=self.noise_map,
             grids=self.grids,
             transformer=self.dataset.transformer,
-            w_tilde=self.w_tilde,
+            sparse_operator=self.dataset.sparse_operator,
         )
 
         return GalaxiesToInversion(

autogalaxy/profiles/mass/dark/mcr_util.py

@@ -146,6 +146,8 @@ def kappa_s_and_scale_radius_for_ludlow(
     if isinstance(mass_at_200, (float, np.ndarray, np.float64)):
         xp = np
     else:
+        import jax.numpy as jnp
+
         xp = jnp
 
     # ------------------------------------
@@ -217,6 +219,8 @@ def kappa_s_scale_radius_and_core_radius_for_ludlow(
     if isinstance(mass_at_200, (float, np.ndarray, np.float64)):
         xp = np
     else:
+        import jax.numpy as jnp
+
         xp = jnp
 
         # ------------------------------------

autogalaxy/profiles/mass/stellar/gaussian.py

@@ -192,7 +192,7 @@ def zeta_from(self, grid: aa.type.Grid2DLike, xp=np):
         q = self.axis_ratio(xp)
         q2 = q**2.0
 
-        ind_pos_y = grid.array[:, 0] >=0
+        ind_pos_y = grid.array[:, 0] >= 0
         shape_grid = np.shape(grid)
         output_grid = np.zeros((shape_grid[0]), dtype=np.complex128)
 
@@ -208,18 +208,24 @@ def zeta_from(self, grid: aa.type.Grid2DLike, xp=np):
         z1_1 = xs_1 + 1j * ys_1
         z2_1 = q * xs_1 + 1j * ys_1 / q
 
-        exp_term_0 = xp.exp(-(xs_0 ** 2) * (1.0 - q2) - ys_0 ** 2 * (1.0 / q2 - 1.0))
-        exp_term_1 = xp.exp(-(xs_1 ** 2) * (1.0 - q2) - ys_1 ** 2 * (1.0 / q2 - 1.0))
+        exp_term_0 = xp.exp(-(xs_0**2) * (1.0 - q2) - ys_0**2 * (1.0 / q2 - 1.0))
+        exp_term_1 = xp.exp(-(xs_1**2) * (1.0 - q2) - ys_1**2 * (1.0 / q2 - 1.0))
 
         if xp == np:
             from scipy.special import wofz
 
             output_grid[ind_pos_y] = -1j * (wofz(z1_0) - exp_term_0 * wofz(z2_0))
-            output_grid[~ind_pos_y] = xp.conj(-1j * (wofz(z1_1) - exp_term_1 * wofz(z2_1)))
+            output_grid[~ind_pos_y] = xp.conj(
+                -1j * (wofz(z1_1) - exp_term_1 * wofz(z2_1))
+            )
 
         else:
-            output_grid[ind_pos_y] = -1j * (self.wofz(z1_0, xp=xp) - exp_term_0 * self.wofz(z2_0, xp=xp))
-            output_grid[~ind_pos_y] = xp.conj(-1j * (self.wofz(z1_1, xp=xp) - exp_term_1 * self.wofz(z2_1, xp=xp)))
+            output_grid[ind_pos_y] = -1j * (
+                self.wofz(z1_0, xp=xp) - exp_term_0 * self.wofz(z2_0, xp=xp)
+            )
+            output_grid[~ind_pos_y] = xp.conj(
+                -1j * (self.wofz(z1_1, xp=xp) - exp_term_1 * self.wofz(z2_1, xp=xp))
+            )
 
         return output_grid
 

test_autogalaxy/profiles/mass/stellar/test_gaussian.py

@@ -241,6 +241,7 @@ def test__image_2d_via_radii_from__correct_value():
 
     assert intensity == pytest.approx(0.32465, 1e-2)
 
+
 def test__wofz():
     from scipy.special import wofz
 
@@ -263,4 +264,3 @@ def test__wofz():
     assert wofz_approx_reg_1 == pytest.approx(wofz(7.0 + 1j * 0.1), 1e-4)
     assert wofz_approx_reg_2 == pytest.approx(wofz(7.0 + 1j * 1e-11), 1e-4)
     assert wofz_approx_reg_3 == pytest.approx(wofz(2.0 + 1j * 1.0), 1e-4)
-