Merge pull request #182 from Jammy2211/feature/jax_speed_up_general

Feature/jax speed up general

Author: Jammy2211

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