Jax to Main!

.github/workflows/main.yml

@@ -22,7 +22,7 @@ jobs:
       uses: actions/setup-python@v4
       with:
         python-version: ${{ matrix.python-version }}
-    - uses: actions/cache@v2
+    - uses: actions/cache@v3
       id: cache-pip
       with:
         path: ~/.cache/pip
@@ -36,9 +36,9 @@ jobs:
         pip3 install setuptools
         pip3 install wheel
         pip3 install pytest coverage pytest-cov
-        pip3 install -r PyAutoConf/requirements.txt
-        pip3 install -r PyAutoArray/requirements.txt
-        pip3 install -r PyAutoArray/optional_requirements.txt
+        pip install ./PyAutoConf
+        pip install ./PyAutoArray
+        pip install ./PyAutoArray[optional]
 
         cd PyAutoArray/autoarray/util/nn/src/nn
         export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/home/runner/work/PyAutoArray/PyAutoArray/PyAutoArray/autoarray/util/nn/src/nn

autoarray/__init__.py

@@ -1,9 +1,13 @@
+from autoconf.dictable import register_parser
+from autoconf 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
@@ -38,20 +42,21 @@
 from .inversion.pixelization.mappers.rectangular import MapperRectangular
 from .inversion.pixelization.mappers.delaunay import MapperDelaunay
 from .inversion.pixelization.mappers.voronoi import MapperVoronoi
+from .inversion.pixelization.mappers.rectangular_uniform import MapperRectangularUniform
 from .inversion.pixelization.image_mesh.abstract import AbstractImageMesh
 from .inversion.pixelization.mesh.abstract import AbstractMesh
 from .inversion.inversion.imaging.mapping import InversionImagingMapping
 from .inversion.inversion.imaging.w_tilde import InversionImagingWTilde
 from .inversion.inversion.interferometer.w_tilde import InversionInterferometerWTilde
 from .inversion.inversion.interferometer.mapping import InversionInterferometerMapping
-from .inversion.inversion.interferometer.lop import InversionInterferometerMappingPyLops
 from .inversion.linear_obj.linear_obj import LinearObj
 from .inversion.linear_obj.func_list import AbstractLinearObjFuncList
 from .mask.derive.indexes_2d import DeriveIndexes2D
 from .mask.derive.mask_1d import DeriveMask1D
 from .mask.derive.mask_2d import DeriveMask2D
 from .mask.derive.grid_1d import DeriveGrid1D
 from .mask.derive.grid_2d import DeriveGrid2D
+from .mask.derive.zoom_2d import Zoom2D
 from .mask.mask_1d import Mask1D
 from .mask.mask_2d import Mask2D
 from .operators.transformer import TransformerDFT
@@ -60,14 +65,17 @@
 from .operators.contour import Grid2DContour
 from .layout.layout import Layout1D
 from .layout.layout import Layout2D
+from .preloads import Preloads
 from .structures.arrays.uniform_1d import Array1D
 from .structures.arrays.uniform_2d import Array2D
+from .structures.arrays.rgb import Array2DRGB
 from .structures.arrays.irregular import ArrayIrregular
 from .structures.grids.uniform_1d import Grid1D
 from .structures.grids.uniform_2d import Grid2D
 from .operators.over_sampling.over_sampler import OverSampler
 from .structures.grids.irregular_2d import Grid2DIrregular
 from .structures.mesh.rectangular_2d import Mesh2DRectangular
+from .structures.mesh.rectangular_2d_uniform import Mesh2DRectangularUniform
 from .structures.mesh.voronoi_2d import Mesh2DVoronoi
 from .structures.mesh.delaunay_2d import Mesh2DDelaunay
 from .structures.arrays.kernel_2d import Kernel2D

autoarray/abstract_ndarray.py

@@ -282,7 +282,7 @@ def output_to_fits(self, file_path: str, overwrite: bool = False):
             If a file already exists at the path, if overwrite=True it is overwritten else an error is raised.
         """
         output_to_fits(
-            values=self.native.array,
+            values=self.native.array.astype("float"),
             file_path=file_path,
             overwrite=overwrite,
             header_dict=self.mask.header_dict,

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/config/visualize/general.yaml

@@ -4,7 +4,6 @@ general:
   log10_min_value: 1.0e-4               # If negative values are being plotted on a log10 scale, values below this value are rounded up to it (e.g. to remove negative values).
   log10_max_value: 1.0e99               # If positive values are being plotted on a log10 scale, values above this value are rounded down to it (e.g. to prevent white blobs).
   zoom_around_mask: true                # If True, plots of data structures with a mask automatically zoom in the masked region.
-  disable_zoom_for_fits: true           # If True, the zoom-in around the masked region is disabled when outputting .fits files, which is useful to retain the same dimensions as the input data.
 inversion:
   reconstruction_vmax_factor: 0.5
   total_mappings_pixels : 8                 # The number of source pixels used when plotting the subplot_mappings of a pixelization.

autoarray/dataset/grids.py

@@ -3,11 +3,11 @@
 from autoarray.mask.mask_2d import Mask2D
 from autoarray.structures.arrays.uniform_2d import Array2D
 from autoarray.structures.arrays.kernel_2d import Kernel2D
-from autoarray.structures.grids.uniform_1d import Grid1D
 from autoarray.structures.grids.uniform_2d import Grid2D
 
 from autoarray.inversion.pixelization.border_relocator import BorderRelocator
-from autoconf import cached_property
+
+from autoarray import exc
 
 
 class GridsDataset:
@@ -24,7 +24,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 +60,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

@@ -15,7 +15,7 @@
 
 from autoarray import exc
 from autoarray.operators.over_sampling import over_sample_util
-from autoarray.inversion.inversion.imaging import inversion_imaging_util
+from autoarray.inversion.inversion.imaging import inversion_imaging_numba_util
 
 logger = logging.getLogger(__name__)
 
@@ -159,9 +159,26 @@ def __init__(
                     """
                 )
 
-        if psf is not None and use_normalized_psf:
+        if psf is not None:
+
+            if not data.mask.is_all_false:
+
+                image_mask = data.mask
+                blurring_mask = data.mask.derive_mask.blurring_from(
+                    kernel_shape_native=psf.shape_native
+                )
+
+            else:
+
+                image_mask = None
+                blurring_mask = None
+
             psf = Kernel2D.no_mask(
-                values=psf.native._array, pixel_scales=psf.pixel_scales, normalize=True
+                values=psf.native._array,
+                pixel_scales=psf.pixel_scales,
+                normalize=use_normalized_psf,
+                image_mask=image_mask,
+                blurring_mask=blurring_mask,
             )
 
         self.psf = psf
@@ -170,9 +187,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,
@@ -203,17 +218,22 @@ def w_tilde(self):
             curvature_preload,
             indexes,
             lengths,
-        ) = inversion_imaging_util.w_tilde_curvature_preload_imaging_from(
+        ) = inversion_imaging_numba_util.w_tilde_curvature_preload_imaging_from(
             noise_map_native=np.array(self.noise_map.native.array).astype("float64"),
             kernel_native=np.array(self.psf.native.array).astype("float64"),
-            native_index_for_slim_index=np.array(self.mask.derive_indexes.native_for_slim).astype("int"),
+            native_index_for_slim_index=np.array(
+                self.mask.derive_indexes.native_for_slim
+            ).astype("int"),
         )
 
         return WTildeImaging(
             curvature_preload=curvature_preload,
             indexes=indexes.astype("int"),
             lengths=lengths.astype("int"),
             noise_map_value=self.noise_map[0],
+            noise_map=self.noise_map,
+            psf=self.psf,
+            mask=self.mask,
         )
 
     @classmethod
@@ -409,12 +429,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,
             )
@@ -488,7 +508,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,
@@ -499,6 +519,8 @@ def apply_over_sampling(
             check_noise_map=False,
         )
 
+        return dataset
+
     def output_to_fits(
         self,
         data_path: Union[Path, str],

autoarray/dataset/imaging/simulator.py

@@ -151,7 +151,7 @@ def via_image_from(
                 pixel_scales=image.pixel_scales,
             )
 
-        if np.isnan(noise_map).any():
+        if np.isnan(noise_map.array).any():
             raise exc.DatasetException(
                 "The noise-map has NaN values in it. This suggests your exposure time and / or"
                 "background sky levels are too low, creating signal counts at or close to 0.0."
@@ -161,7 +161,9 @@ def via_image_from(
             image = image - background_sky_map
 
         mask = Mask2D.all_false(
-            shape_native=image.shape_native, pixel_scales=image.pixel_scales
+            shape_native=image.shape_native,
+            pixel_scales=image.pixel_scales,
+            origin=image.origin,
         )
 
         image = Array2D(values=image, mask=mask)