Merge pull request #179 from Jammy2211/feature/jax_remove_profiling

Feature/jax remove profiling

Author: Jammy2211

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