Feature/jax remove profiling

autoarray/__init__.py

@@ -1,9 +1,13 @@
+from autoconf.dictable import register_parser
+from autofit import conf
+
+conf.instance.register(__file__)
+
 from . import exc
 from . import type
 from . import util
 from . import fixtures
 from . import mock as m
-from .numba_util import profile_func
 from .dataset import preprocess
 from .dataset.abstract.dataset import AbstractDataset
 from .dataset.abstract.w_tilde import AbstractWTilde

autoarray/config/general.yaml

@@ -1,3 +1,5 @@
+jax:
+  use_jax: true                       # If True, uses JAX internally, whereas False uses normal Numpy.
 fits:
   flip_for_ds9: false                 # If True, the image is flipped before output to a .fits file, which is useful for viewing in DS9.
 inversion:

autoarray/dataset/grids.py

@@ -9,6 +9,8 @@
 from autoarray.inversion.pixelization.border_relocator import BorderRelocator
 from autoconf import cached_property
 
+from autoarray import exc
+
 
 class GridsDataset:
     def __init__(
@@ -24,7 +26,7 @@ def __init__(
 
         The following grids are contained:
 
-        - `uniform`: A grids of (y,x) coordinates which aligns with the centre of every image pixel of the image data,
+        - `lp`: A grids of (y,x) coordinates which aligns with the centre of every image pixel of the image data,
         which is used for most normal calculations (e.g. evaluating the amount of light that falls in an pixel
         from a light profile).
 
@@ -60,72 +62,30 @@ def __init__(
         self.over_sample_size_pixelization = over_sample_size_pixelization
         self.psf = psf
 
-    @cached_property
-    def lp(self) -> Union[Grid1D, Grid2D]:
-        """
-        Returns the grid of (y,x) Cartesian coordinates at the centre of every pixel in the masked data, which is used
-        to perform most normal calculations (e.g. evaluating the amount of light that falls in an pixel from a light
-        profile).
-
-        This grid is computed based on the mask, in particular its pixel-scale and sub-grid size.
-
-        Returns
-        -------
-        The (y,x) coordinates of every pixel in the data.
-        """
-        return Grid2D.from_mask(
+        self.lp = Grid2D.from_mask(
             mask=self.mask,
             over_sample_size=self.over_sample_size_lp,
         )
+        self.lp.over_sampled
 
-    @cached_property
-    def pixelization(self) -> Grid2D:
-        """
-        Returns the grid of (y,x) Cartesian coordinates of every pixel in the masked data which is used
-        specifically for calculations associated with a pixelization.
-
-        The `pixelization` grid is identical to the `uniform` grid but often uses a different over sampling scheme
-        when performing calculations. For example, the pixelization may benefit from using a a higher `sub_size` than
-        the `uniform` grid, in order to better prevent aliasing effects.
-
-        This grid is computed based on the mask, in particular its pixel-scale and sub-grid size.
-
-        Returns
-        -------
-        The (y,x) coordinates of every pixel in the data, used for pixelization / inversion calculations.
-        """
-        return Grid2D.from_mask(
+        self.pixelization = Grid2D.from_mask(
             mask=self.mask,
             over_sample_size=self.over_sample_size_pixelization,
         )
-
-    @cached_property
-    def blurring(self) -> Optional[Grid2D]:
-        """
-        Returns a blurring-grid from a mask and the 2D shape of the PSF kernel.
-
-        A blurring grid consists of all pixels that are masked (and therefore have their values set to (0.0, 0.0)),
-        but are close enough to the unmasked pixels that their values will be convolved into the unmasked those pixels.
-        This when computing images from light profile objects.
-
-        This uses lazy allocation such that the calculation is only performed when the blurring grid is used, ensuring
-        efficient set up of the `Imaging` class.
-
-        Returns
-        -------
-        The blurring grid given the mask of the imaging data.
-        """
+        self.pixelization.over_sampled
 
         if self.psf is None:
-            return None
-
-        return self.lp.blurring_grid_via_kernel_shape_from(
-            kernel_shape_native=self.psf.shape_native,
-        )
-
-    @cached_property
-    def border_relocator(self) -> BorderRelocator:
-        return BorderRelocator(
+            self.blurring = None
+        else:
+            try:
+                self.blurring = self.lp.blurring_grid_via_kernel_shape_from(
+                    kernel_shape_native=self.psf.shape_native,
+                )
+                self.blurring.over_sampled
+            except exc.MaskException:
+                self.blurring = None
+
+        self.border_relocator = BorderRelocator(
             mask=self.mask, sub_size=self.over_sample_size_pixelization
         )
 

autoarray/dataset/imaging/dataset.py

@@ -170,9 +170,7 @@ def __init__(
             if psf.mask.shape[0] % 2 == 0 or psf.mask.shape[1] % 2 == 0:
                 raise exc.KernelException("Kernel2D Kernel2D must be odd")
 
-    @cached_property
-    def grids(self):
-        return GridsDataset(
+        self.grids = GridsDataset(
             mask=self.data.mask,
             over_sample_size_lp=self.over_sample_size_lp,
             over_sample_size_pixelization=self.over_sample_size_pixelization,
@@ -511,7 +509,7 @@ def apply_over_sampling(
             passed into the calculations performed in the `inversion` module.
         """
 
-        return Imaging(
+        dataset = Imaging(
             data=self.data,
             noise_map=self.noise_map,
             psf=self.psf,
@@ -522,6 +520,8 @@ def apply_over_sampling(
             check_noise_map=False,
         )
 
+        return dataset
+
     def output_to_fits(
         self,
         data_path: Union[Path, str],

autoarray/dataset/interferometer/dataset.py

@@ -101,9 +101,7 @@ def __init__(
             else None
         )
 
-    @cached_property
-    def grids(self):
-        return GridsDataset(
+        self.grids = GridsDataset(
             mask=self.real_space_mask,
             over_sample_size_lp=self.over_sample_size_lp,
             over_sample_size_pixelization=self.over_sample_size_pixelization,

autoarray/exc.py

@@ -104,11 +104,3 @@ class PlottingException(Exception):
     """
 
     pass
-
-
-class ProfilingException(Exception):
-    """
-    Raises exceptions associated with in-built profiling tools (e.g. the `profile_func` decorator).
-    """
-
-    pass

autoarray/fit/fit_dataset.py

@@ -13,7 +13,6 @@
 from autoarray.inversion.inversion.abstract import AbstractInversion
 from autoarray.mask.mask_2d import Mask2D
 
-from autoarray.numba_util import profile_func
 from autoarray import type as ty
 
 
@@ -116,7 +115,6 @@ def __init__(
         dataset,
         use_mask_in_fit: bool = False,
         dataset_model: DatasetModel = None,
-        run_time_dict: Optional[Dict] = None,
     ):
         """Class to fit a masked dataset where the dataset's data structures are any dimension.
 
@@ -149,7 +147,6 @@ def __init__(
         self.dataset = dataset
         self.use_mask_in_fit = use_mask_in_fit
         self.dataset_model = dataset_model or DatasetModel()
-        self.run_time_dict = run_time_dict
 
     @property
     def mask(self) -> Mask2D:
@@ -320,7 +317,6 @@ def log_evidence(self) -> float:
             )
 
     @property
-    @profile_func
     def figure_of_merit(self) -> float:
         if self.inversion is not None:
             return self.log_evidence

autoarray/fit/fit_imaging.py

@@ -14,7 +14,6 @@ def __init__(
         dataset: Imaging,
         use_mask_in_fit: bool = False,
         dataset_model: DatasetModel = None,
-        run_time_dict: Optional[Dict] = None,
     ):
         """
         Class to fit a masked imaging dataset.
@@ -50,7 +49,6 @@ def __init__(
             dataset=dataset,
             use_mask_in_fit=use_mask_in_fit,
             dataset_model=dataset_model,
-            run_time_dict=run_time_dict,
         )
 
     @property

autoarray/fit/fit_interferometer.py

@@ -18,7 +18,6 @@ def __init__(
         dataset: Interferometer,
         dataset_model: DatasetModel = None,
         use_mask_in_fit: bool = False,
-        run_time_dict: Optional[Dict] = None,
     ):
         """
         Class to fit a masked interferometer dataset.
@@ -59,7 +58,6 @@ def __init__(
             dataset=dataset,
             dataset_model=dataset_model,
             use_mask_in_fit=use_mask_in_fit,
-            run_time_dict=run_time_dict,
         )
 
     @property

autoarray/fit/mock/mock_fit_imaging.py

@@ -15,13 +15,11 @@ def __init__(
         model_data=None,
         inversion=None,
         blurred_image=None,
-        run_time_dict: Optional[Dict] = None,
     ):
         super().__init__(
             dataset=dataset or MockDataset(),
             dataset_model=dataset_model,
             use_mask_in_fit=use_mask_in_fit,
-            run_time_dict=run_time_dict,
         )
 
         self._noise_map = noise_map

autoarray/geometry/geometry_util.py

@@ -2,8 +2,6 @@
 import numpy as np
 from typing import Tuple, Union
 
-
-from autoarray import numba_util
 from autoarray import type as ty
 
 
@@ -178,70 +176,6 @@ def convert_pixel_scales_2d(pixel_scales: ty.PixelScales) -> Tuple[float, float]
     return pixel_scales
 
 
-@numba_util.jit()
-def central_pixel_coordinates_2d_numba_from(
-    shape_native: Tuple[int, int],
-) -> Tuple[float, float]:
-    """
-    Returns the central pixel coordinates of a 2D geometry (and therefore a 2D data structure like an ``Array2D``)
-    from the shape of that data structure.
-
-    Examples of the central pixels are as follows:
-
-    - For a 3x3 image, the central pixel is pixel [1, 1].
-    - For a 4x4 image, the central pixel is [1.5, 1.5].
-
-    Parameters
-    ----------
-    shape_native
-        The dimensions of the data structure, which can be in 1D, 2D or higher dimensions.
-
-    Returns
-    -------
-    The central pixel coordinates of the data structure.
-    """
-    return (float(shape_native[0] - 1) / 2, float(shape_native[1] - 1) / 2)
-
-
-@numba_util.jit()
-def central_scaled_coordinate_2d_numba_from(
-    shape_native: Tuple[int, int],
-    pixel_scales: ty.PixelScales,
-    origin: Tuple[float, float] = (0.0, 0.0),
-) -> Tuple[float, float]:
-    """
-    Returns the central scaled coordinates of a 2D geometry (and therefore a 2D data structure like an ``Array2D``)
-    from the shape of that data structure.
-
-    This is computed by using the data structure's shape and converting it to scaled units using an input
-    pixel-coordinates to scaled-coordinate conversion factor `pixel_scales`.
-
-    The origin of the scaled grid can also be input and moved from (0.0, 0.0).
-
-    Parameters
-    ----------
-    shape_native
-        The 2D shape of the data structure whose central scaled coordinates are computed.
-    pixel_scales
-        The (y,x) scaled units to pixel units conversion factor of the 2D data structure.
-    origin
-        The (y,x) scaled units origin of the coordinate system the central scaled coordinate is computed on.
-
-    Returns
-    -------
-    The central coordinates of the 2D data structure in scaled units.
-    """
-
-    central_pixel_coordinates = central_pixel_coordinates_2d_numba_from(
-        shape_native=shape_native
-    )
-
-    y_pixel = central_pixel_coordinates[0] + (origin[0] / pixel_scales[0])
-    x_pixel = central_pixel_coordinates[1] - (origin[1] / pixel_scales[1])
-
-    return (y_pixel, x_pixel)
-
-
 def central_pixel_coordinates_2d_from(
     shape_native: Tuple[int, int],
 ) -> Tuple[float, float]:
@@ -294,7 +228,7 @@ def central_scaled_coordinate_2d_from(
     The central coordinates of the 2D data structure in scaled units.
     """
 
-    central_pixel_coordinates = central_pixel_coordinates_2d_numba_from(
+    central_pixel_coordinates = central_pixel_coordinates_2d_from(
         shape_native=shape_native
     )
 
@@ -367,7 +301,6 @@ def pixel_coordinates_2d_from(
     return (y_pixel, x_pixel)
 
 
-@numba_util.jit()
 def scaled_coordinates_2d_from(
     pixel_coordinates_2d: Tuple[float, float],
     shape_native: Tuple[int, int],
@@ -411,7 +344,7 @@ def scaled_coordinates_2d_from(
         origin=(0.0, 0.0)
     )
     """
-    central_scaled_coordinates = central_scaled_coordinate_2d_numba_from(
+    central_scaled_coordinates = central_scaled_coordinate_2d_from(
         shape_native=shape_native, pixel_scales=pixel_scales, origin=origins
     )