Merge pull request #157 from Jammy2211/feature/jax_merge

feature/jax merge

Author: Jammy2211

autoarray/__init__.py

@@ -14,7 +14,6 @@
 from .dataset.interferometer.dataset import Interferometer
 from .dataset.interferometer.simulator import SimulatorInterferometer
 from .dataset.interferometer.w_tilde import WTildeInterferometer
-from .dataset.over_sampling import OverSamplingDataset
 from .dataset.dataset_model import DatasetModel
 from .fit.fit_dataset import AbstractFit
 from .fit.fit_dataset import FitDataset
@@ -68,14 +67,8 @@
 from .structures.arrays.irregular import ArrayIrregular
 from .structures.grids.uniform_1d import Grid1D
 from .structures.grids.uniform_2d import Grid2D
-from .operators.over_sampling.decorator import perform_over_sampling_from
-from .operators.over_sampling.grid_oversampled import Grid2DOverSampled
-from .operators.over_sampling.uniform import OverSamplingUniform
-from .operators.over_sampling.iterate import OverSamplingIterate
-from .operators.over_sampling.uniform import OverSamplerUniform
-from .operators.over_sampling.iterate import OverSamplerIterate
+from .operators.over_sampling.over_sampler import OverSampler
 from .structures.grids.irregular_2d import Grid2DIrregular
-from .structures.grids.irregular_2d import Grid2DIrregularUniform
 from .structures.mesh.rectangular_2d import Mesh2DRectangular
 from .structures.mesh.voronoi_2d import Mesh2DVoronoi
 from .structures.mesh.delaunay_2d import Mesh2DDelaunay
@@ -98,4 +91,4 @@
 
 conf.instance.register(__file__)
 
-__version__ = "2024.11.6.1"
+__version__ = "2025.1.18.7"

autoarray/config/visualize/mat_wrap_2d.yaml

@@ -146,6 +146,15 @@ VectorYXQuiver:      # wrapper for `plt.quiver()`: customize (y,x) vectors appea
     linewidth: 5
     pivot: middle
     units: xy
+DelaunayDrawer:       # wrapper for `plt.fill()`: customize the appearance of Delaunay mesh's.
+  figure:
+    alpha: 0.7
+    edgecolor: k
+    linewidth: 0.0
+  subplot:
+    alpha: 0.7
+    edgecolor: k
+    linewidth: 0.0  
 VoronoiDrawer:       # wrapper for `plt.fill()`: customize the appearance of Voronoi mesh's.
   figure:
     alpha: 0.7

autoarray/config/visualize/plots.yaml

@@ -5,20 +5,11 @@
 
 dataset:                                   # Settings for plots of all datasets (e.g. ImagingPlotter, InterferometerPlotter).
   subplot_dataset: true                    # Plot subplot containing all dataset quantities (e.g. the data, noise-map, etc.)?
-  data: false                              # Plot the individual data of every dataset?
-  noise_map: false                         # Plot the individual noise-map of every dataset?
-  signal_to_noise_map: false               # Plot the individual signal-to-noise-map of every dataset?
-  over_sampling: false            # Plot the over-sampling sub-size, used to evaluate light profiles, of every dataset?
-  over_sampling_non_uniform: false  # Plot the over-sampling sub-size, used to evaluate non uniform grids, of every dataset?
-  over_sampling_pixelization: false  # Plot the over-sampling sub-size, used to evaluate pixelizations, of every dataset?
 imaging:                                   # Settings for plots of imaging datasets (e.g. ImagingPlotter)
    psf: false
 fit:                                       # Settings for plots of all fits (e.g. FitImagingPlotter, FitInterferometerPlotter).
   subplot_fit: true                        # Plot subplot of all fit quantities for any dataset (e.g. the model data, residual-map, etc.)?
   subplot_fit_log10: true                  # Plot subplot of all fit quantities for any dataset using log10 color maps (e.g. the model data, residual-map, etc.)?
-  all_at_end_png: true                     # Plot all individual plots listed below as .png (even if False)?
-  all_at_end_fits: true                    # Plot all individual plots listed below as .fits (even if False)?
-  all_at_end_pdf: false                    # Plot all individual plots listed below as publication-quality .pdf (even if False)?
   data: false                              # Plot individual plots of the data?
   noise_map: false                         # Plot individual plots of the noise-map?
   signal_to_noise_map: false               # Plot individual plots of the signal-to-noise-map?
@@ -31,33 +22,14 @@ fit_imaging: {}                            # Settings for plots of fits to imagi
 inversion:                                 # Settings for plots of inversions (e.g. InversionPlotter).
   subplot_inversion: true                  # Plot subplot of all quantities in each inversion (e.g. reconstrucuted image, reconstruction)?
   subplot_mappings: true                   # Plot subplot of the image-to-source pixels mappings of each pixelization?
-  all_at_end_png: true                     # Plot all individual plots listed below as .png (even if False)?
-  all_at_end_fits: true                    # Plot all individual plots listed below as .fits (even if False)?
-  all_at_end_pdf: false                    # Plot all individual plots listed below as publication-quality .pdf (even if False)?
   data_subtracted: false                   # Plot individual plots of the data with the other inversion linear objects subtracted?
-  errors: false                            # Plot image of the errors of every mesh-pixel reconstructed value?
+  reconstruction_noise_map: false          # Plot image of the noise of every mesh-pixel reconstructed value?
   sub_pixels_per_image_pixels: false       # Plot the number of sub pixels per masked data pixels?
   mesh_pixels_per_image_pixels: false      # Plot the number of image-plane mesh pixels per masked data pixels?
   image_pixels_per_mesh_pixels: false      # Plot the number of image pixels in each pixel of the mesh?
   reconstructed_image: false               # Plot image of the reconstructed data (e.g. in the image-plane)?
   reconstruction: false                    # Plot the reconstructed inversion (e.g. the pixelization's mesh in the source-plane)?
   regularization_weights: false            # Plot the effective regularization weight of every inversion mesh pixel?
-interferometer:                            # Settings for plots of interferometer datasets (e.g. InterferometerPlotter).
-  amplitudes_vs_uv_distances: false
-  phases_vs_uv_distances: false
-  uv_wavelengths: false
-  dirty_image: false
-  dirty_noise_map: false
-  dirty_signal_to_noise_map: false
 fit_interferometer:                        # Settings for plots of fits to interferometer datasets (e.g. FitInterferometerPlotter).
   subplot_fit_dirty_images: false          # Plot subplot of the dirty-images of all interferometer datasets?
-  subplot_fit_real_space: false            # Plot subplot of the real-space images of all interferometer datasets?
-  amplitudes_vs_uv_distances: false
-  phases_vs_uv_distances: false
-  uv_wavelengths: false
-  dirty_image: false
-  dirty_noise_map: false
-  dirty_signal_to_noise_map: false
-  dirty_residual_map: false
-  dirty_normalized_residual_map: false
-  dirty_chi_squared_map: false
+  subplot_fit_real_space: false            # Plot subplot of the real-space images of all interferometer datasets?

autoarray/dataset/abstract/dataset.py

@@ -4,15 +4,17 @@
 import warnings
 from typing import Optional, Union
 
-from autoarray.dataset.over_sampling import OverSamplingDataset
+from autoconf import cached_property
+
 from autoarray.dataset.grids import GridsDataset
 
 from autoarray import exc
 from autoarray.mask.mask_1d import Mask1D
 from autoarray.mask.mask_2d import Mask2D
 from autoarray.structures.abstract_structure import Structure
 from autoarray.structures.arrays.uniform_2d import Array2D
-from autoconf import cached_property
+
+from autoarray.operators.over_sampling import over_sample_util
 
 
 logger = logging.getLogger(__name__)
@@ -24,7 +26,8 @@ def __init__(
         data: Structure,
         noise_map: Structure,
         noise_covariance_matrix: Optional[np.ndarray] = None,
-        over_sampling: Optional[OverSamplingDataset] = OverSamplingDataset(),
+        over_sample_size_lp: Union[int, Array2D] = 4,
+        over_sample_size_pixelization: Union[int, Array2D] = 4,
     ):
         """
         An abstract dataset, containing the image data, noise-map, PSF and associated quantities for calculations
@@ -45,6 +48,32 @@ def __init__(
         over sampling calculations built in which approximate the 2D line integral of these calculations within a
         pixel. This is explained in more detail in the `GridsDataset` class.
 
+        **Over Sampling**
+
+        If a grid is uniform and the centre of each point on the grid is the centre of a 2D pixel, evaluating
+        the value of a function on the grid requires a 2D line integral to compute it precisely. This can be
+        computationally expensive and difficult to implement.
+
+        Over sampling is a numerical technique where the function is evaluated on a sub-grid within each grid pixel
+        which is higher resolution than the grid itself. This approximates more closely the value of the function
+        within a 2D line intergral of the values in the square pixel that the grid is centred.
+
+        For example, in PyAutoGalaxy and PyAutoLens the light profiles and galaxies are evaluated in order to determine
+        how much light falls in each pixel. This uses over sampling and therefore a higher resolution grid than the
+        image data to ensure the calculation is accurate.
+
+        This class controls how over sampling is performed for 2 different types of grids:
+
+        - `lp`: A grids of (y,x) coordinates which aligns with the centre of every image pixel of the image data
+        and is used to evaluate light profiles for model-fititng.
+
+        - `pixelization`: A grid of (y,x) coordinates which again align with the centre of every image pixel of
+        the image data. This grid is used specifically for pixelizations computed via the `inversion` module, which
+        can benefit from using different oversampling schemes than the normal grid.
+
+        Different calculations typically benefit from different over sampling, which this class enables
+        the customization of.
+
         Parameters
         ----------
         data
@@ -56,10 +85,13 @@ def __init__(
         noise_covariance_matrix
             A noise-map covariance matrix representing the covariance between noise in every `data` value, which
             can be used via a bespoke fit to account for correlated noise in the data.
-        over_sampling
-            The over sampling schemes which divide the grids into sub grids of smaller pixels within their host image
-            pixels when using the grid to evaluate a function (e.g. images) to better approximate the 2D line integral
-            This class controls over sampling for all the different grids (e.g. `grid`, `grids.pixelization).
+        over_sample_size_lp
+            The over sampling scheme size, which divides the grid into a sub grid of smaller pixels when computing
+            values (e.g. images) from the grid to approximate the 2D line integral of the amount of light that falls
+            into each pixel.
+        over_sample_size_pixelization
+            How over sampling is performed for the grid which is associated with a pixelization, which is therefore
+            passed into the calculations performed in the `inversion` module.
         """
 
         self.data = data
@@ -93,15 +125,28 @@ def __init__(
 
         self.noise_map = noise_map
 
-        self.over_sampling = over_sampling
+        self.over_sample_size_lp = (
+            over_sample_util.over_sample_size_convert_to_array_2d_from(
+                over_sample_size=over_sample_size_lp, mask=self.mask
+            )
+        )
+        self.over_sample_size_pixelization = (
+            over_sample_util.over_sample_size_convert_to_array_2d_from(
+                over_sample_size=over_sample_size_pixelization, mask=self.mask
+            )
+        )
 
     @property
     def grid(self):
-        return self.grids.uniform
+        return self.grids.lp
 
     @cached_property
     def grids(self):
-        return GridsDataset(mask=self.data.mask, over_sampling=self.over_sampling)
+        return GridsDataset(
+            mask=self.data.mask,
+            over_sample_size_lp=self.over_sample_size_lp,
+            over_sample_size_pixelization=self.over_sample_size_pixelization,
+        )
 
     @property
     def shape_native(self):
@@ -161,4 +206,16 @@ def trimmed_after_convolution_from(self, kernel_shape) -> "AbstractDataset":
             kernel_shape=kernel_shape
         )
 
+        dataset.over_sample_size_lp = (
+            dataset.over_sample_size_lp.trimmed_after_convolution_from(
+                kernel_shape=kernel_shape
+            )
+        )
+
+        dataset.over_sample_size_pixelization = (
+            dataset.over_sample_size_pixelization.trimmed_after_convolution_from(
+                kernel_shape=kernel_shape
+            )
+        )
+
         return dataset

autoarray/dataset/grids.py

@@ -1,12 +1,11 @@
 from typing import Optional, Union
 
-from autoarray.dataset.over_sampling import OverSamplingDataset
 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.operators.over_sampling.uniform import OverSamplingUniform
 from autoarray.inversion.pixelization.border_relocator import BorderRelocator
 from autoconf import cached_property
 
@@ -15,7 +14,8 @@ class GridsDataset:
     def __init__(
         self,
         mask: Mask2D,
-        over_sampling: OverSamplingDataset,
+        over_sample_size_lp: Union[int, Array2D],
+        over_sample_size_pixelization: Union[int, Array2D],
         psf: Optional[Kernel2D] = None,
     ):
         """
@@ -28,11 +28,6 @@ def __init__(
         which is used for most normal calculations (e.g. evaluating the amount of light that falls in an pixel
         from a light profile).
 
-        - `non_uniform`: A grid of (y,x) coordinates which aligns with the centre of every image pixel of the image
-        data, but where their values are going to be deflected to become non-uniform such that the adaptive over
-        sampling scheme used for the main grid does not apply. This is used to compute over sampled light profiles of
-        lensed sources in PyAutoLens.
-
         - `pixelization`: A grid of (y,x) coordinates which again align with the centre of every image pixel of
         the image data. This grid is used specifically for pixelizations computed via the `inversion` module, which
         can benefit from using different oversampling schemes than the normal grid.
@@ -49,16 +44,24 @@ def __init__(
 
         Parameters
         ----------
-        mask
-        over_sampling
+        over_sample_size_lp
+            The over sampling scheme size, which divides the grid into a sub grid of smaller pixels when computing
+            values (e.g. images) from the grid to approximate the 2D line integral of the amount of light that falls
+            into each pixel.
+        over_sample_size_pixelization
+            How over sampling is performed for the grid which is associated with a pixelization, which is therefore
+            passed into the calculations performed in the `inversion` module.
         psf
+            The Point Spread Function kernel of the image which accounts for diffraction due to the telescope optics
+            via 2D convolution.
         """
         self.mask = mask
-        self.over_sampling = over_sampling
+        self.over_sample_size_lp = over_sample_size_lp
+        self.over_sample_size_pixelization = over_sample_size_pixelization
         self.psf = psf
 
     @cached_property
-    def uniform(self) -> Union[Grid1D, Grid2D]:
+    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
@@ -70,38 +73,9 @@ def uniform(self) -> Union[Grid1D, Grid2D]:
         -------
         The (y,x) coordinates of every pixel in the data.
         """
-
-        return Grid2D.from_mask(
-            mask=self.mask,
-            over_sampling=self.over_sampling.uniform,
-        )
-
-    @cached_property
-    def non_uniform(self) -> Optional[Union[Grid1D, Grid2D]]:
-        """
-        Returns the grid of (y,x) Cartesian coordinates at the centre of every pixel in the masked data, but
-        with a different over sampling scheme designed for
-
-        where
-        their values are going to be deflected to become non-uniform such that the adaptive over sampling scheme used
-        for the main grid does not apply.
-
-        This is used to compute over sampled light profiles of lensed sources in PyAutoLens.
-
-
-        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.
-        """
-
-        if self.over_sampling.non_uniform is None:
-            return None
-
         return Grid2D.from_mask(
             mask=self.mask,
-            over_sampling=self.over_sampling.non_uniform,
+            over_sample_size=self.over_sample_size_lp,
         )
 
     @cached_property
@@ -120,15 +94,9 @@ def pixelization(self) -> Grid2D:
         -------
         The (y,x) coordinates of every pixel in the data, used for pixelization / inversion calculations.
         """
-
-        over_sampling = self.over_sampling.pixelization
-
-        if over_sampling is None:
-            over_sampling = OverSamplingUniform(sub_size=4)
-
         return Grid2D.from_mask(
             mask=self.mask,
-            over_sampling=over_sampling,
+            over_sample_size=self.over_sample_size_pixelization,
         )
 
     @cached_property
@@ -151,36 +119,26 @@ def blurring(self) -> Optional[Grid2D]:
         if self.psf is None:
             return None
 
-        return self.uniform.blurring_grid_via_kernel_shape_from(
+        return self.lp.blurring_grid_via_kernel_shape_from(
             kernel_shape_native=self.psf.shape_native,
         )
 
-    @cached_property
-    def over_sampler_non_uniform(self):
-        return self.non_uniform.over_sampling.over_sampler_from(mask=self.mask)
-
-    @cached_property
-    def over_sampler_pixelization(self):
-        return self.pixelization.over_sampling.over_sampler_from(mask=self.mask)
-
     @cached_property
     def border_relocator(self) -> BorderRelocator:
         return BorderRelocator(
-            mask=self.mask, sub_size=self.pixelization.over_sampling.sub_size
+            mask=self.mask, sub_size=self.over_sample_size_pixelization
         )
 
 
 class GridsInterface:
     def __init__(
         self,
-        uniform=None,
-        non_uniform=None,
+        lp=None,
         pixelization=None,
         blurring=None,
         border_relocator=None,
     ):
-        self.uniform = uniform
-        self.non_uniform = non_uniform
+        self.lp = lp
         self.pixelization = pixelization
         self.blurring = blurring
         self.border_relocator = border_relocator