Feature/jax speed up general

autoarray/__init__.py

@@ -1,5 +1,5 @@
 from autoconf.dictable import register_parser
-from autofit import conf
+from autoconf import conf
 
 conf.instance.register(__file__)
 

autoarray/dataset/imaging/dataset.py

@@ -430,12 +430,12 @@ def apply_noise_scaling(
         """
 
         if signal_to_noise_value is None:
-            noise_map = np.array(self.noise_map.native.array)
+            noise_map = self.noise_map.native
             noise_map[mask.array == False] = noise_value
         else:
             noise_map = np.where(
                 mask == False,
-                np.median(self.data.native.array[mask.derive_mask.edge == False])
+                np.median(self.data.native[mask.derive_mask.edge == False])
                 / signal_to_noise_value,
                 self.noise_map.native.array,
             )

autoarray/dataset/interferometer/dataset.py

@@ -1,4 +1,3 @@
-from astropy.io import fits
 import logging
 import numpy as np
 from pathlib import Path
@@ -148,6 +147,9 @@ def from_fits(
         )
 
     def w_tilde_preprocessing(self):
+
+        from astropy.io import fits
+
         if self.preprocessing_directory.is_dir():
             filename = "{}/curvature_preload.fits".format(self.preprocessing_directory)
 

autoarray/dataset/plot/imaging_plotters.py

@@ -14,9 +14,9 @@ def __init__(
         self,
         dataset: Imaging,
         get_visuals_2d: Callable,
-        mat_plot_2d: MatPlot2D = MatPlot2D(),
-        visuals_2d: Visuals2D = Visuals2D(),
-        include_2d: Include2D = Include2D(),
+        mat_plot_2d: MatPlot2D = None,
+        visuals_2d: Visuals2D = None,
+        include_2d: Include2D = None,
     ):
         """
         Plots the attributes of `Imaging` objects using the matplotlib method `imshow()` and many other matplotlib
@@ -231,9 +231,9 @@ class ImagingPlotter(Plotter):
     def __init__(
         self,
         dataset: Imaging,
-        mat_plot_2d: MatPlot2D = MatPlot2D(),
-        visuals_2d: Visuals2D = Visuals2D(),
-        include_2d: Include2D = Include2D(),
+        mat_plot_2d: MatPlot2D = None,
+        visuals_2d: Visuals2D = None,
+        include_2d: Include2D = None,
     ):
         """
         Plots the attributes of `Imaging` objects using the matplotlib method `imshow()` and many other matplotlib

autoarray/dataset/plot/interferometer_plotters.py

@@ -14,12 +14,12 @@ class InterferometerPlotter(Plotter):
     def __init__(
         self,
         dataset: Interferometer,
-        mat_plot_1d: MatPlot1D = MatPlot1D(),
-        visuals_1d: Visuals1D = Visuals1D(),
-        include_1d: Include1D = Include1D(),
-        mat_plot_2d: MatPlot2D = MatPlot2D(),
-        visuals_2d: Visuals2D = Visuals2D(),
-        include_2d: Include2D = Include2D(),
+        mat_plot_1d: MatPlot1D = None,
+        visuals_1d: Visuals1D = None,
+        include_1d: Include1D = None,
+        mat_plot_2d: MatPlot2D = None,
+        visuals_2d: Visuals2D = None,
+        include_2d: Include2D = None,
     ):
         """
         Plots the attributes of `Interferometer` objects using the matplotlib methods `plot()`, `scatter()` and

autoarray/dataset/preprocess.py

@@ -1,5 +1,4 @@
 import numpy as np
-from scipy.stats import norm
 
 from autoarray import exc
 
@@ -316,6 +315,7 @@ def background_noise_map_via_edges_from(image, no_edges):
     no_edges
         Number of edges used to estimate the background level.
     """
+    from scipy.stats import norm
 
     from autoarray.structures.arrays.uniform_2d import Array2D
 

autoarray/fit/fit_interferometer.py

@@ -122,7 +122,7 @@ def noise_normalization(self) -> float:
         [Noise_Term] = sum(log(2*pi*[Noise]**2.0))
         """
         return fit_util.noise_normalization_complex_from(
-            noise_map=self.noise_map,
+            noise_map=self.noise_map.array,
         )
 
     @property

autoarray/fit/fit_util.py

@@ -174,12 +174,8 @@ def noise_normalization_complex_from(*, noise_map: jnp.ndarray) -> float:
     noise_map
         The masked noise-map of the dataset.
     """
-    noise_normalization_real = jnp.sum(
-        jnp.log(2 * jnp.pi * np.array(noise_map).real ** 2.0)
-    )
-    noise_normalization_imag = jnp.sum(
-        jnp.log(2 * jnp.pi * np.array(noise_map).imag ** 2.0)
-    )
+    noise_normalization_real = jnp.sum(jnp.log(2 * jnp.pi * noise_map.real**2.0))
+    noise_normalization_imag = jnp.sum(jnp.log(2 * jnp.pi * noise_map.imag**2.0))
     return noise_normalization_real + noise_normalization_imag
 
 

autoarray/fit/mock/mock_fit_imaging.py

@@ -1,4 +1,4 @@
-from typing import Dict, Optional
+from typing import Optional
 
 from autoarray.dataset.mock.mock_dataset import MockDataset
 from autoarray.dataset.dataset_model import DatasetModel

autoarray/fit/plot/fit_imaging_plotters.py

@@ -13,9 +13,9 @@ def __init__(
         self,
         fit,
         get_visuals_2d: Callable,
-        mat_plot_2d: MatPlot2D = MatPlot2D(),
-        visuals_2d: Visuals2D = Visuals2D(),
-        include_2d: Include2D = Include2D(),
+        mat_plot_2d: MatPlot2D = None,
+        visuals_2d: Visuals2D = None,
+        include_2d: Include2D = None,
         residuals_symmetric_cmap: bool = True,
     ):
         """
@@ -242,9 +242,9 @@ class FitImagingPlotter(Plotter):
     def __init__(
         self,
         fit: FitImaging,
-        mat_plot_2d: MatPlot2D = MatPlot2D(),
-        visuals_2d: Visuals2D = Visuals2D(),
-        include_2d: Include2D = Include2D(),
+        mat_plot_2d: MatPlot2D = None,
+        visuals_2d: Visuals2D = None,
+        include_2d: Include2D = None,
     ):
         """
         Plots the attributes of `FitImaging` objects using the matplotlib method `imshow()` and many other matplotlib

autoarray/fit/plot/fit_interferometer_plotters.py

@@ -20,9 +20,9 @@ def __init__(
         mat_plot_1d: MatPlot1D,
         visuals_1d: Visuals1D,
         include_1d: Include1D,
-        mat_plot_2d: MatPlot2D = MatPlot2D(),
-        visuals_2d: Visuals2D = Visuals2D(),
-        include_2d: Include2D = Include2D(),
+        mat_plot_2d: MatPlot2D = None,
+        visuals_2d: Visuals2D = None,
+        include_2d: Include2D = None,
         residuals_symmetric_cmap: bool = True,
     ):
         """
@@ -455,12 +455,12 @@ class FitInterferometerPlotter(Plotter):
     def __init__(
         self,
         fit: FitInterferometer,
-        mat_plot_1d: MatPlot1D = MatPlot1D(),
-        visuals_1d: Visuals1D = Visuals1D(),
-        include_1d: Include1D = Include1D(),
-        mat_plot_2d: MatPlot2D = MatPlot2D(),
-        visuals_2d: Visuals2D = Visuals2D(),
-        include_2d: Include2D = Include2D(),
+        mat_plot_1d: MatPlot1D = None,
+        visuals_1d: Visuals1D = None,
+        include_1d: Include1D = None,
+        mat_plot_2d: MatPlot2D = None,
+        visuals_2d: Visuals2D = None,
+        include_2d: Include2D = None,
     ):
         """
         Plots the attributes of `FitInterferometer` objects using the matplotlib method `imshow()` and many other

autoarray/fit/plot/fit_vector_yx_plotters.py

@@ -14,9 +14,9 @@ def __init__(
         self,
         fit,
         get_visuals_2d: Callable,
-        mat_plot_2d: MatPlot2D = MatPlot2D(),
-        visuals_2d: Visuals2D = Visuals2D(),
-        include_2d: Include2D = Include2D(),
+        mat_plot_2d: MatPlot2D = None,
+        visuals_2d: Visuals2D = None,
+        include_2d: Include2D = None,
     ):
         """
         Plots the attributes of `FitImaging` objects using the matplotlib method `imshow()` and many other matplotlib
@@ -208,9 +208,9 @@ class FitImagingPlotter(Plotter):
     def __init__(
         self,
         fit: FitImaging,
-        mat_plot_2d: MatPlot2D = MatPlot2D(),
-        visuals_2d: Visuals2D = Visuals2D(),
-        include_2d: Include2D = Include2D(),
+        mat_plot_2d: MatPlot2D = None,
+        visuals_2d: Visuals2D = None,
+        include_2d: Include2D = None,
     ):
         """
         Plots the attributes of `FitImaging` objects using the matplotlib method `imshow()` and many other matplotlib

autoarray/geometry/geometry_2d.py

@@ -154,7 +154,7 @@ def pixel_coordinates_2d_from(
         A 2D (y,x) pixel-value coordinate.
         """
         return geometry_util.pixel_coordinates_2d_from(
-            scaled_coordinates_2d=np.array(scaled_coordinates_2d),
+            scaled_coordinates_2d=scaled_coordinates_2d,
             shape_native=self.shape_native,
             pixel_scales=self.pixel_scales,
             origins=self.origin,
@@ -184,7 +184,7 @@ def scaled_coordinates_2d_from(
         """
 
         return geometry_util.scaled_coordinates_2d_from(
-            pixel_coordinates_2d=np.array(pixel_coordinates_2d),
+            pixel_coordinates_2d=pixel_coordinates_2d,
             shape_native=self.shape_native,
             pixel_scales=self.pixel_scales,
             origins=self.origin,
@@ -235,7 +235,7 @@ def grid_pixels_2d_from(self, grid_scaled_2d: Grid2D) -> Grid2D:
         from autoarray.structures.grids.uniform_2d import Grid2D
 
         grid_pixels_2d = geometry_util.grid_pixels_2d_slim_from(
-            grid_scaled_2d_slim=np.array(grid_scaled_2d.array),
+            grid_scaled_2d_slim=grid_scaled_2d.array,
             shape_native=self.shape_native,
             pixel_scales=self.pixel_scales,
             origin=self.origin,
@@ -261,7 +261,7 @@ def grid_pixel_centres_2d_from(self, grid_scaled_2d: Grid2D) -> Grid2D:
         from autoarray.structures.grids.uniform_2d import Grid2D
 
         grid_pixel_centres_1d = geometry_util.grid_pixel_centres_2d_slim_from(
-            grid_scaled_2d_slim=np.array(grid_scaled_2d),
+            grid_scaled_2d_slim=grid_scaled_2d,
             shape_native=self.shape_native,
             pixel_scales=self.pixel_scales,
             origin=self.origin,
@@ -294,7 +294,7 @@ def grid_pixel_indexes_2d_from(self, grid_scaled_2d: Grid2D) -> Array2D:
         from autoarray.structures.arrays.uniform_2d import Array2D
 
         grid_pixel_indexes_2d = geometry_util.grid_pixel_indexes_2d_slim_from(
-            grid_scaled_2d_slim=np.array(grid_scaled_2d),
+            grid_scaled_2d_slim=grid_scaled_2d,
             shape_native=self.shape_native,
             pixel_scales=self.pixel_scales,
             origin=self.origin,
@@ -320,7 +320,7 @@ def grid_scaled_2d_from(self, grid_pixels_2d: Grid2D) -> Grid2D:
         from autoarray.structures.grids.uniform_2d import Grid2D
 
         grid_scaled_1d = geometry_util.grid_scaled_2d_slim_from(
-            grid_pixels_2d_slim=np.array(grid_pixels_2d),
+            grid_pixels_2d_slim=grid_pixels_2d,
             shape_native=self.shape_native,
             pixel_scales=self.pixel_scales,
             origin=self.origin,

autoarray/geometry/geometry_util.py

@@ -474,8 +474,8 @@ def grid_pixels_2d_slim_from(
     centres_scaled = central_scaled_coordinate_2d_from(
         shape_native=shape_native, pixel_scales=pixel_scales, origin=origin
     )
-    centres_scaled = np.array(centres_scaled)
-    pixel_scales = np.array(pixel_scales)
+    centres_scaled = centres_scaled
+    pixel_scales = pixel_scales
     sign = np.array([-1, 1])
     return (sign * grid_scaled_2d_slim / pixel_scales) + centres_scaled + 0.5
 

autoarray/inversion/inversion/abstract.py

@@ -1,17 +1,13 @@
 import copy
-import jax
 import jax.numpy as jnp
 import numpy as np
-from scipy.linalg import block_diag
-from scipy.sparse import csc_matrix
-from scipy.sparse.linalg import splu
+
 from typing import Dict, List, Optional, Type, Union
 
 from autoconf import cached_property
 
 from autoarray.dataset.imaging.dataset import Imaging
 from autoarray.dataset.interferometer.dataset import Interferometer
-from autoarray.inversion.inversion.mapper_valued import MapperValued
 from autoarray.inversion.inversion.dataset_interface import DatasetInterface
 from autoarray.inversion.linear_obj.linear_obj import LinearObj
 from autoarray.inversion.pixelization.mappers.abstract import AbstractMapper
@@ -337,6 +333,8 @@ def regularization_matrix(self) -> Optional[np.ndarray]:
         If the `settings.force_edge_pixels_to_zeros` is `True`, the edge pixels of each mapper in the inversion
         are regularized so high their value is forced to zero.
         """
+        from scipy.linalg import block_diag
+
         return block_diag(
             *[linear_obj.regularization_matrix for linear_obj in self.linear_obj_list]
         )
@@ -703,6 +701,8 @@ def log_det_regularization_matrix_term(self) -> float:
         float
             The log determinant of the regularization matrix.
         """
+        from scipy.sparse import csc_matrix
+        from scipy.sparse.linalg import splu
 
         if not self.has(cls=AbstractRegularization):
             return 0.0

autoarray/inversion/inversion/imaging/inversion_imaging_util.py

@@ -1,11 +1,7 @@
 import numpy as np
-from scipy.linalg import cho_solve
-from typing import List, Optional, Tuple
-
-from autoarray.inversion.inversion.settings import SettingsInversion
+from typing import Tuple
 
 from autoarray import numba_util
-from autoarray import exc
 
 
 @numba_util.jit()

autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py

@@ -1,4 +1,3 @@
-from astropy.io import fits
 import logging
 import numpy as np
 import time
@@ -791,6 +790,9 @@ def w_tilde_curvature_preload_interferometer_in_stages_with_chunks_from(
     check=True,
     directory=None,
 ) -> np.ndarray:
+
+    from astropy.io import fits
+
     if directory is None:
         raise NotImplementedError()
 

autoarray/inversion/inversion/inversion_util.py

@@ -1,10 +1,8 @@
 import jax.numpy as jnp
-import jaxnnls
-import jax
 import jax.lax as lax
 import numpy as np
 
-from typing import List, Optional, Tuple
+from typing import List, Optional
 
 from autoconf import conf
 
@@ -302,6 +300,7 @@ def reconstruction_positive_only_from(
     -------
     Non-negative S that minimizes the Eq.(2) of https://arxiv.org/pdf/astro-ph/0302587.pdf.
     """
+    import jaxnnls
 
     try:
         reconstruction = jaxnnls.solve_nnls_primal(curvature_reg_matrix, data_vector)

autoarray/inversion/mock/mock_inversion.py

@@ -29,13 +29,15 @@ def __init__(
         regularization_term=None,
         log_det_curvature_reg_matrix_term=None,
         log_det_regularization_matrix_term=None,
-        settings: SettingsInversion = SettingsInversion(),
+        settings: SettingsInversion = None,
     ):
         dataset = DatasetInterface(
             data=data,
             noise_map=noise_map,
         )
 
+        settings = settings or SettingsInversion()
+
         super().__init__(
             dataset=dataset,
             linear_obj_list=linear_obj_list or [],