From be571b5a3b849ba545fd19aadf89fed2a4765dfc Mon Sep 17 00:00:00 2001
From: James Nightingale <james.w.nightingale@durham.ac.uk>
Date: Fri, 20 Dec 2024 11:21:02 +0000
Subject: [PATCH] remove OverSamplingDataset

---
 autoarray/__init__.py                         |  1 -
 autoarray/dataset/abstract/dataset.py         | 63 ++++++++++++---
 autoarray/dataset/grids.py                    | 25 ++++--
 autoarray/dataset/imaging/dataset.py          | 77 +++++++++++--------
 autoarray/dataset/imaging/simulator.py        | 12 ++-
 autoarray/dataset/interferometer/dataset.py   | 15 ++--
 autoarray/dataset/over_sampling.py            | 52 -------------
 autoarray/dataset/plot/imaging_plotters.py    | 28 +++----
 autoarray/fixtures.py                         | 10 +--
 .../pixelization/border_relocator.py          |  9 +--
 .../pixelization/image_mesh/hilbert.py        |  1 +
 .../pixelization/mappers/delaunay.py          |  8 +-
 .../over_sampling/over_sample_util.py         | 39 +++++++++-
 .../operators/over_sampling/over_sampler.py   |  9 +--
 autoarray/structures/grids/irregular_2d.py    |  2 +-
 autoarray/structures/grids/uniform_2d.py      | 21 +++--
 .../structures/triangles/array/__init__.py    |  1 +
 .../triangles/coordinate_array/__init__.py    |  1 +
 .../dataset/abstract/test_dataset.py          | 36 ++++-----
 .../dataset/plot/test_imaging_plotters.py     |  8 +-
 .../inversion/plot/test_mapper_plotters.py    | 11 +--
 test_autoarray/inversion/test_linear_obj.py   |  4 +-
 test_autoarray/plot/get_visuals/test_two_d.py |  8 +-
 .../structures/grids/test_uniform_2d.py       |  7 +-
 24 files changed, 242 insertions(+), 206 deletions(-)
 delete mode 100644 autoarray/dataset/over_sampling.py

diff --git a/autoarray/__init__.py b/autoarray/__init__.py
index 0e102360..58a32da2 100644
--- a/autoarray/__init__.py
+++ b/autoarray/__init__.py
@@ -15,7 +15,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
diff --git a/autoarray/dataset/abstract/dataset.py b/autoarray/dataset/abstract/dataset.py
index ea7c1671..eca0d7c8 100644
--- a/autoarray/dataset/abstract/dataset.py
+++ b/autoarray/dataset/abstract/dataset.py
@@ -4,7 +4,8 @@
 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
@@ -12,7 +13,8 @@
 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,7 +125,16 @@ 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):
@@ -101,7 +142,11 @@ def grid(self):
 
     @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):
diff --git a/autoarray/dataset/grids.py b/autoarray/dataset/grids.py
index 95231b70..c460bf82 100644
--- a/autoarray/dataset/grids.py
+++ b/autoarray/dataset/grids.py
@@ -1,7 +1,7 @@
 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
@@ -14,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,
     ):
         """
@@ -43,12 +44,20 @@ 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
@@ -66,7 +75,7 @@ def lp(self) -> Union[Grid1D, Grid2D]:
         """
         return Grid2D.from_mask(
             mask=self.mask,
-            over_sample_size=self.over_sampling.lp,
+            over_sample_size=self.over_sample_size_lp,
         )
 
     @cached_property
@@ -87,7 +96,7 @@ def pixelization(self) -> Grid2D:
         """
         return Grid2D.from_mask(
             mask=self.mask,
-            over_sample_size=self.over_sampling.pixelization,
+            over_sample_size=self.over_sample_size_pixelization,
         )
 
     @cached_property
@@ -117,7 +126,7 @@ def blurring(self) -> Optional[Grid2D]:
     @cached_property
     def border_relocator(self) -> BorderRelocator:
         return BorderRelocator(
-            mask=self.mask, sub_size=self.pixelization.over_sample_size
+            mask=self.mask, sub_size=self.over_sample_size_pixelization
         )
 
 
diff --git a/autoarray/dataset/imaging/dataset.py b/autoarray/dataset/imaging/dataset.py
index 2c3e50b8..938a812d 100644
--- a/autoarray/dataset/imaging/dataset.py
+++ b/autoarray/dataset/imaging/dataset.py
@@ -8,7 +8,6 @@
 from autoarray.dataset.abstract.dataset import AbstractDataset
 from autoarray.dataset.grids import GridsDataset
 from autoarray.dataset.imaging.w_tilde import WTildeImaging
-from autoarray.dataset.over_sampling import OverSamplingDataset
 from autoarray.structures.arrays.uniform_2d import Array2D
 from autoarray.operators.convolver import Convolver
 from autoarray.structures.arrays.kernel_2d import Kernel2D
@@ -28,7 +27,8 @@ def __init__(
         noise_map: Optional[Array2D] = None,
         psf: Optional[Kernel2D] = None,
         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,
         pad_for_convolver: bool = False,
         use_normalized_psf: Optional[bool] = True,
         check_noise_map: bool = True,
@@ -69,10 +69,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.
         pad_for_convolver
             The PSF convolution may extend beyond the edges of the image mask, which can lead to edge effects in the
             convolved image. If `True`, the image and noise-map are padded to ensure the PSF convolution does not
@@ -114,7 +117,8 @@ def __init__(
             data=data,
             noise_map=noise_map,
             noise_covariance_matrix=noise_covariance_matrix,
-            over_sampling=over_sampling,
+            over_sample_size_lp=over_sample_size_lp,
+            over_sample_size_pixelization=over_sample_size_pixelization,
         )
 
         self.use_normalized_psf = use_normalized_psf
@@ -143,7 +147,10 @@ def __init__(
     @cached_property
     def grids(self):
         return GridsDataset(
-            mask=self.data.mask, over_sampling=self.over_sampling, psf=self.psf
+            mask=self.data.mask,
+            over_sample_size_lp=self.over_sample_size_lp,
+            over_sample_size_pixelization=self.over_sample_size_pixelization,
+            psf=self.psf,
         )
 
     @cached_property
@@ -214,7 +221,8 @@ def from_fits(
         psf_hdu: int = 0,
         noise_covariance_matrix: Optional[np.ndarray] = None,
         check_noise_map: bool = True,
-        over_sampling: Optional[OverSamplingDataset] = OverSamplingDataset(),
+        over_sample_size_lp: Union[int, Array2D] = 4,
+        over_sample_size_pixelization: Union[int, Array2D] = 4,
     ) -> "Imaging":
         """
         Load an imaging dataset from multiple .fits file.
@@ -253,10 +261,13 @@ def from_fits(
             A noise-map covariance matrix representing the covariance between noise in every `data` value.
         check_noise_map
             If True, the noise-map is checked to ensure all values are above zero.
-        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.
         """
 
         data = Array2D.from_fits(
@@ -284,7 +295,8 @@ def from_fits(
             psf=psf,
             noise_covariance_matrix=noise_covariance_matrix,
             check_noise_map=check_noise_map,
-            over_sampling=over_sampling,
+            over_sample_size_lp=over_sample_size_lp,
+            over_sample_size_pixelization=over_sample_size_pixelization,
         )
 
     def apply_mask(self, mask: Mask2D) -> "Imaging":
@@ -323,12 +335,18 @@ def apply_mask(self, mask: Mask2D) -> "Imaging":
         else:
             noise_covariance_matrix = None
 
+        over_sample_size_lp = Array2D(values=self.over_sample_size_lp.native, mask=mask)
+        over_sample_size_pixelization = Array2D(
+            values=self.over_sample_size_pixelization.native, mask=mask
+        )
+
         dataset = Imaging(
             data=data,
             noise_map=noise_map,
             psf=self.psf,
             noise_covariance_matrix=noise_covariance_matrix,
-            over_sampling=self.over_sampling,
+            over_sample_size_lp=over_sample_size_lp,
+            over_sample_size_pixelization=over_sample_size_pixelization,
             pad_for_convolver=True,
         )
 
@@ -416,7 +434,8 @@ def apply_noise_scaling(
             noise_map=noise_map,
             psf=self.psf,
             noise_covariance_matrix=self.noise_covariance_matrix,
-            over_sampling=self.over_sampling,
+            over_sample_size_lp=self.over_sample_size_lp,
+            over_sample_size_pixelization=self.over_sample_size_pixelization,
             pad_for_convolver=False,
             check_noise_map=False,
         )
@@ -429,7 +448,8 @@ def apply_noise_scaling(
 
     def apply_over_sampling(
         self,
-        over_sampling: Optional[OverSamplingDataset] = OverSamplingDataset(),
+        over_sample_size_lp: Union[int, Array2D] = None,
+        over_sample_size_pixelization: Union[int, Array2D] = None,
     ) -> "AbstractDataset":
         """
         Apply new over sampling objects to the grid and grid pixelization of the dataset.
@@ -449,25 +469,22 @@ def apply_over_sampling(
 
         Parameters
         ----------
-        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.
         """
 
-        uniform = over_sampling.lp or self.over_sampling.lp
-        pixelization = over_sampling.pixelization or self.over_sampling.pixelization
-
-        over_sampling = OverSamplingDataset(
-            lp=uniform,
-            pixelization=pixelization,
-        )
-
         return Imaging(
             data=self.data,
             noise_map=self.noise_map,
             psf=self.psf,
-            over_sampling=over_sampling,
+            over_sample_size_lp=over_sample_size_lp or self.over_sample_size_lp,
+            over_sample_size_pixelization=over_sample_size_pixelization
+            or self.over_sample_size_pixelization,
             pad_for_convolver=False,
             check_noise_map=False,
         )
diff --git a/autoarray/dataset/imaging/simulator.py b/autoarray/dataset/imaging/simulator.py
index 47f4d4da..a43cc614 100644
--- a/autoarray/dataset/imaging/simulator.py
+++ b/autoarray/dataset/imaging/simulator.py
@@ -1,5 +1,6 @@
 import logging
 import numpy as np
+from typing import Optional, Union
 
 from autoarray.dataset.imaging.dataset import Imaging
 from autoarray.structures.arrays.uniform_2d import Array2D
@@ -98,7 +99,9 @@ def __init__(
         self.noise_if_add_noise_false = noise_if_add_noise_false
         self.noise_seed = noise_seed
 
-    def via_image_from(self, image: Array2D) -> Imaging:
+    def via_image_from(
+        self, image: Array2D, over_sample_size: Optional[Union[int, np.ndarray]] = None
+    ) -> Imaging:
         """
         Simulate an `Imaging` dataset from an input image.
 
@@ -163,6 +166,11 @@ def via_image_from(self, image: Array2D) -> Imaging:
 
         image = Array2D(values=image, mask=mask)
 
-        return Imaging(
+        dataset = Imaging(
             data=image, psf=self.psf, noise_map=noise_map, check_noise_map=False
         )
+
+        if over_sample_size is not None:
+            dataset = dataset.apply_over_sampling(over_sample_size_lp=over_sample_size)
+
+        return dataset
diff --git a/autoarray/dataset/interferometer/dataset.py b/autoarray/dataset/interferometer/dataset.py
index abc155fd..f69753dd 100644
--- a/autoarray/dataset/interferometer/dataset.py
+++ b/autoarray/dataset/interferometer/dataset.py
@@ -7,7 +7,6 @@
 from autoarray.dataset.abstract.dataset import AbstractDataset
 from autoarray.dataset.interferometer.w_tilde import WTildeInterferometer
 from autoarray.dataset.grids import GridsDataset
-from autoarray.dataset.over_sampling import OverSamplingDataset
 from autoarray.operators.transformer import TransformerNUFFT
 
 from autoarray.structures.visibilities import Visibilities
@@ -74,15 +73,11 @@ def __init__(
         """
         self.real_space_mask = real_space_mask
 
-        over_sampling = OverSamplingDataset(
-            lp=1,
-            pixelization=1,
-        )
-
         super().__init__(
             data=data,
             noise_map=noise_map,
-            over_sampling=over_sampling,
+            over_sample_size_lp=1,
+            over_sample_size_pixelization=1,
         )
 
         self.uv_wavelengths = uv_wavelengths
@@ -93,7 +88,11 @@ def __init__(
 
     @cached_property
     def grids(self):
-        return GridsDataset(mask=self.real_space_mask, over_sampling=self.over_sampling)
+        return GridsDataset(
+            mask=self.real_space_mask,
+            over_sample_size_lp=self.over_sample_size_lp,
+            over_sample_size_pixelization=self.over_sample_size_pixelization,
+        )
 
     @classmethod
     def from_fits(
diff --git a/autoarray/dataset/over_sampling.py b/autoarray/dataset/over_sampling.py
deleted file mode 100644
index c19938e3..00000000
--- a/autoarray/dataset/over_sampling.py
+++ /dev/null
@@ -1,52 +0,0 @@
-import logging
-from typing import Union
-
-from autoarray.structures.arrays.uniform_2d import Array2D
-
-logger = logging.getLogger(__name__)
-
-
-class OverSamplingDataset:
-    def __init__(
-        self,
-        lp: Union[int, Array2D] = 4,
-        pixelization: Union[int, Array2D] = 4,
-    ):
-        """
-        Customizes how over sampling calculations are performed using the grids of the data.
-
-        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 3 different types of grids:
-
-        - `grid`: A grids of (y,x) coordinates which aligns with the centre of every image pixel of the image data.
-
-        - `grids.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
-        ----------
-        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 so as to approximate the 2D line integral of the amount of light that falls
-            into each pixel.
-        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.lp = lp
-        self.pixelization = pixelization
diff --git a/autoarray/dataset/plot/imaging_plotters.py b/autoarray/dataset/plot/imaging_plotters.py
index 3f8bac92..8fed6f60 100644
--- a/autoarray/dataset/plot/imaging_plotters.py
+++ b/autoarray/dataset/plot/imaging_plotters.py
@@ -61,8 +61,8 @@ def figures_2d(
         noise_map: bool = False,
         psf: bool = False,
         signal_to_noise_map: bool = False,
-        over_sampling: bool = False,
-        over_sampling_pixelization: bool = False,
+        over_sample_size_lp: bool = False,
+        over_sample_size_pixelization: bool = False,
         title_str: Optional[str] = None,
     ):
         """
@@ -81,10 +81,10 @@ def figures_2d(
             Whether to make a 2D plot (via `imshow`) of the psf.
         signal_to_noise_map
             Whether to make a 2D plot (via `imshow`) of the signal-to-noise map.
-        over_sampling
+        over_sample_size_lp
             Whether to make a 2D plot (via `imshow`) of the Over Sampling for input light profiles. If
             adaptive sub size is used, the sub size grid for a centre of (0.0, 0.0) is used.
-        over_sampling_pixelization
+        over_sample_size_pixelization
             Whether to make a 2D plot (via `imshow`) of the Over Sampling for pixelizations.
         """
 
@@ -124,24 +124,24 @@ def figures_2d(
                 ),
             )
 
-        if over_sampling:
+        if over_sample_size_lp:
             self.mat_plot_2d.plot_array(
-                array=self.dataset.grids.lp.over_sample_size,
+                array=self.dataset.grids.over_sample_size_lp,
                 visuals_2d=self.get_visuals_2d(),
                 auto_labels=AutoLabels(
-                    title=title_str or f"Over Sampling (Uniform)",
-                    filename="over_sampling",
+                    title=title_str or f"Over Sample Size (Light Profiles)",
+                    filename="over_sample_size_lp",
                     cb_unit="",
                 ),
             )
 
-        if over_sampling_pixelization:
+        if over_sample_size_pixelization:
             self.mat_plot_2d.plot_array(
-                array=self.dataset.grids.pixelization.over_sample_size,
+                array=self.dataset.grids.over_sample_size_pixelization,
                 visuals_2d=self.get_visuals_2d(),
                 auto_labels=AutoLabels(
-                    title=title_str or f"Over Sampling (Pixelization)",
-                    filename="over_sampling_pixelization",
+                    title=title_str or f"Over Sample Size (Pixelization)",
+                    filename="over_sample_size_pixelization",
                     cb_unit="",
                 ),
             )
@@ -218,8 +218,8 @@ def subplot_dataset(self):
 
         self.figures_2d(signal_to_noise_map=True)
 
-        self.figures_2d(over_sampling=True)
-        self.figures_2d(over_sampling_pixelization=True)
+        self.figures_2d(over_sample_size_lp=True)
+        self.figures_2d(over_sample_size_pixelization=True)
 
         self.mat_plot_2d.output.subplot_to_figure(auto_filename="subplot_dataset")
         self.close_subplot_figure()
diff --git a/autoarray/fixtures.py b/autoarray/fixtures.py
index 73d150d4..acfd277d 100644
--- a/autoarray/fixtures.py
+++ b/autoarray/fixtures.py
@@ -160,7 +160,7 @@ def make_imaging_7x7():
         data=make_image_7x7(),
         psf=make_psf_3x3(),
         noise_map=make_noise_map_7x7(),
-        over_sampling=aa.OverSamplingDataset(lp=1),
+        over_sample_size_lp=1,
     )
 
 
@@ -169,7 +169,7 @@ def make_imaging_7x7_sub_2():
         data=make_image_7x7(),
         psf=make_psf_3x3(),
         noise_map=make_noise_map_7x7(),
-        over_sampling=aa.OverSamplingDataset(lp=2),
+        over_sample_size_lp=2,
     )
 
 
@@ -178,7 +178,7 @@ def make_imaging_covariance_7x7():
         data=make_image_7x7(),
         psf=make_psf_3x3(),
         noise_covariance_matrix=make_noise_covariance_matrix_7x7(),
-        over_sampling=aa.OverSamplingDataset(lp=1),
+        over_sample_size_lp=1,
     )
 
 
@@ -187,7 +187,7 @@ def make_imaging_7x7_no_blur():
         data=make_image_7x7(),
         psf=make_psf_3x3_no_blur(),
         noise_map=make_noise_map_7x7(),
-        over_sampling=aa.OverSamplingDataset(lp=1),
+        over_sample_size_lp=1,
     )
 
 
@@ -196,7 +196,7 @@ def make_imaging_7x7_no_blur_sub_2():
         data=make_image_7x7(),
         psf=make_psf_3x3_no_blur(),
         noise_map=make_noise_map_7x7(),
-        over_sampling=aa.OverSamplingDataset(lp=2),
+        over_sample_size_lp=2,
     )
 
 
diff --git a/autoarray/inversion/pixelization/border_relocator.py b/autoarray/inversion/pixelization/border_relocator.py
index 28f7a3ff..da6d2812 100644
--- a/autoarray/inversion/pixelization/border_relocator.py
+++ b/autoarray/inversion/pixelization/border_relocator.py
@@ -132,12 +132,9 @@ class BorderRelocator:
     def __init__(self, mask: Mask2D, sub_size: Union[int, Array2D]):
         self.mask = mask
 
-        if isinstance(sub_size, int):
-            sub_size = Array2D(
-                values=np.full(fill_value=sub_size, shape=mask.shape_slim), mask=mask
-            )
-
-        self.sub_size = sub_size
+        self.sub_size = over_sample_util.over_sample_size_convert_to_array_2d_from(
+            over_sample_size=sub_size, mask=mask
+        )
 
     @cached_property
     def border_slim(self):
diff --git a/autoarray/inversion/pixelization/image_mesh/hilbert.py b/autoarray/inversion/pixelization/image_mesh/hilbert.py
index bde49849..964457de 100644
--- a/autoarray/inversion/pixelization/image_mesh/hilbert.py
+++ b/autoarray/inversion/pixelization/image_mesh/hilbert.py
@@ -86,6 +86,7 @@ def generate2d(x, y, ax, ay, bx, by):
             -(ay - ay2),
         )
 
+
 def grid_hilbert_order_from(length, mask_radius):
     """
     This function will create a grid in the Hilbert space-filling curve order.
diff --git a/autoarray/inversion/pixelization/mappers/delaunay.py b/autoarray/inversion/pixelization/mappers/delaunay.py
index f18c3b39..737247b0 100644
--- a/autoarray/inversion/pixelization/mappers/delaunay.py
+++ b/autoarray/inversion/pixelization/mappers/delaunay.py
@@ -117,9 +117,7 @@ def pix_sub_weights(self) -> PixSubWeights:
         pix_indexes_for_simplex_index = delaunay.simplices
 
         mappings, sizes = mapper_util.pix_indexes_for_sub_slim_index_delaunay_from(
-            source_plane_data_grid=np.array(
-                self.source_plane_data_grid.over_sampled
-            ),
+            source_plane_data_grid=np.array(self.source_plane_data_grid.over_sampled),
             simplex_index_for_sub_slim_index=simplex_index_for_sub_slim_index,
             pix_indexes_for_simplex_index=pix_indexes_for_simplex_index,
             delaunay_points=delaunay.points,
@@ -129,9 +127,7 @@ def pix_sub_weights(self) -> PixSubWeights:
         sizes = sizes.astype("int")
 
         weights = mapper_util.pixel_weights_delaunay_from(
-            source_plane_data_grid=np.array(
-                self.source_plane_data_grid.over_sampled
-            ),
+            source_plane_data_grid=np.array(self.source_plane_data_grid.over_sampled),
             source_plane_mesh_grid=np.array(self.source_plane_mesh_grid),
             slim_index_for_sub_slim_index=self.slim_index_for_sub_slim_index,
             pix_indexes_for_sub_slim_index=mappings,
diff --git a/autoarray/operators/over_sampling/over_sample_util.py b/autoarray/operators/over_sampling/over_sample_util.py
index c6584c92..1d4637fc 100644
--- a/autoarray/operators/over_sampling/over_sample_util.py
+++ b/autoarray/operators/over_sampling/over_sample_util.py
@@ -1,7 +1,6 @@
 from __future__ import annotations
 import numpy as np
-from typing import TYPE_CHECKING, List, Tuple
-from autoconf import conf
+from typing import TYPE_CHECKING, Union, List, Tuple
 
 from autoarray.structures.arrays.uniform_2d import Array2D
 
@@ -9,13 +8,47 @@
     from autoarray.structures.grids.uniform_2d import Grid2D
 
 from autoarray.geometry import geometry_util
+from autoarray.mask.mask_2d import Mask2D
+
 from autoarray import numba_util
 from autoarray.mask import mask_2d_util
 
-from autoarray import exc
 from autoarray import type as ty
 
 
+def over_sample_size_convert_to_array_2d_from(
+    over_sample_size: Union[int, np.ndarray], mask: Union[np.ndarray, Mask2D]
+):
+    """
+    Returns the over sample size as an `Array2D` object, for example converting it from a single integer.
+
+    The interface allows a user to specify the `over_sample_size` as either:
+
+    - A single integer, whereby over sampling is performed to this degree for every pixel.
+    - An ndarray with the same number of entries as the mask, to enable adaptive over sampling.
+
+    This function converts these input structures to an `Array2D` which is used internally in the source code
+    to perform computations.
+
+    Parameters
+    ----------
+    over_sample_size
+        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.
+
+    Returns
+    -------
+
+    """
+    if isinstance(over_sample_size, int):
+        over_sample_size = np.full(
+            fill_value=over_sample_size, shape=mask.pixels_in_mask
+        ).astype("int")
+
+    return Array2D(values=over_sample_size, mask=mask)
+
+
 @numba_util.jit()
 def total_sub_pixels_2d_from(sub_size: np.ndarray) -> int:
     """
diff --git a/autoarray/operators/over_sampling/over_sampler.py b/autoarray/operators/over_sampling/over_sampler.py
index 1bdb1740..84888baf 100644
--- a/autoarray/operators/over_sampling/over_sampler.py
+++ b/autoarray/operators/over_sampling/over_sampler.py
@@ -141,12 +141,9 @@ def __init__(self, mask: Mask2D, sub_size: Union[int, Array2D]):
         """
         self.mask = mask
 
-        if isinstance(sub_size, int):
-            sub_size = Array2D(
-                values=np.full(fill_value=sub_size, shape=mask.shape_slim), mask=mask
-            )
-
-        self.sub_size = sub_size
+        self.sub_size = over_sample_util.over_sample_size_convert_to_array_2d_from(
+            over_sample_size=sub_size, mask=mask
+        )
 
     def tree_flatten(self):
         return (self.mask,), ()
diff --git a/autoarray/structures/grids/irregular_2d.py b/autoarray/structures/grids/irregular_2d.py
index 54cb483c..2116291b 100644
--- a/autoarray/structures/grids/irregular_2d.py
+++ b/autoarray/structures/grids/irregular_2d.py
@@ -269,4 +269,4 @@ def grid_of_closest_from(self, grid_pair: "Grid2DIrregular") -> "Grid2DIrregular
 
             grid_of_closest[i, :] = self[np.argmin(radial_distances), :]
 
-        return Grid2DIrregular(values=grid_of_closest)
\ No newline at end of file
+        return Grid2DIrregular(values=grid_of_closest)
diff --git a/autoarray/structures/grids/uniform_2d.py b/autoarray/structures/grids/uniform_2d.py
index 604e561a..e4cad24b 100644
--- a/autoarray/structures/grids/uniform_2d.py
+++ b/autoarray/structures/grids/uniform_2d.py
@@ -151,7 +151,7 @@ def __init__(
             mapping large data arrays to and from the slim / native formats, which can be a computational bottleneck.
         over_sample_size
             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 so as to approximate the 2D line integral of the amount of light that falls
+            values (e.g. images) from the grid to approximate the 2D line integral of the amount of light that falls
             into each pixel.
         over_sampled
             The over sampled grid of (y,x) coordinates, which can be passed in manually because if the grid is
@@ -170,12 +170,9 @@ def __init__(
 
         grid_2d_util.check_grid_2d(grid_2d=values)
 
-        if isinstance(over_sample_size, int):
-            over_sample_size = np.full(
-                fill_value=over_sample_size, shape=mask.shape_slim
-            ).astype("int")
-
-        over_sample_size = Array2D(values=over_sample_size, mask=mask)
+        over_sample_size = over_sample_util.over_sample_size_convert_to_array_2d_from(
+            over_sample_size=over_sample_size, mask=mask
+        )
 
         from autoarray.operators.over_sampling.over_sampler import OverSampler
 
@@ -1119,8 +1116,7 @@ def is_uniform(self) -> bool:
         return True
 
     def apply_over_sampling(
-        self,
-        over_sample_size : Union[int, np.ndarray]
+        self, over_sample_size: Union[int, np.ndarray]
     ) -> "AbstractDataset":
         """
         Apply new over sampling to the grid.
@@ -1132,17 +1128,18 @@ def apply_over_sampling(
         ----------
         over_sample_size
             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 so as to approximate the 2D line integral of the amount of light that falls
+            values (e.g. images) from the grid to approximate the 2D line integral of the amount of light that falls
             into each pixel.
         """
         if not self.is_uniform:
             raise exc.GridException(
                 """
                 Cannot apply over sampling to a Grid2D which is not uniform.
-                """)
+                """
+            )
 
         return Grid2D(
             values=self,
             mask=self.mask,
             over_sample_size=over_sample_size,
-        )
\ No newline at end of file
+        )
diff --git a/autoarray/structures/triangles/array/__init__.py b/autoarray/structures/triangles/array/__init__.py
index 58feab94..0fade4b8 100644
--- a/autoarray/structures/triangles/array/__init__.py
+++ b/autoarray/structures/triangles/array/__init__.py
@@ -1,4 +1,5 @@
 from .array import ArrayTriangles
+
 try:
     from .jax_array import ArrayTriangles as JAXArrayTriangles
 except ImportError:
diff --git a/autoarray/structures/triangles/coordinate_array/__init__.py b/autoarray/structures/triangles/coordinate_array/__init__.py
index b40a452b..f70bc8a9 100644
--- a/autoarray/structures/triangles/coordinate_array/__init__.py
+++ b/autoarray/structures/triangles/coordinate_array/__init__.py
@@ -1,4 +1,5 @@
 from .coordinate_array import CoordinateArrayTriangles
+
 try:
     from .jax_coordinate_array import (
         CoordinateArrayTriangles as JAXCoordinateArrayTriangles,
diff --git a/test_autoarray/dataset/abstract/test_dataset.py b/test_autoarray/dataset/abstract/test_dataset.py
index a1dc7573..dd99c120 100644
--- a/test_autoarray/dataset/abstract/test_dataset.py
+++ b/test_autoarray/dataset/abstract/test_dataset.py
@@ -65,7 +65,7 @@ def test__grid__uses_mask_and_settings(
     masked_imaging_7x7 = ds.AbstractDataset(
         data=masked_image_7x7,
         noise_map=masked_noise_map_7x7,
-        over_sampling=aa.OverSamplingDataset(lp=2),
+        over_sample_size_lp=2,
     )
 
     assert isinstance(masked_imaging_7x7.grids.lp, aa.Grid2D)
@@ -94,28 +94,24 @@ def test__grids_pixelization__uses_mask_and_settings(
     masked_imaging_7x7 = ds.AbstractDataset(
         data=masked_image_7x7,
         noise_map=masked_noise_map_7x7,
-        over_sampling=aa.OverSamplingDataset(
-            lp=2,
-            pixelization=4,
-        ),
+        over_sample_size_lp=2,
+        over_sample_size_pixelization=4,
     )
 
     assert isinstance(masked_imaging_7x7.grids.pixelization, aa.Grid2D)
-    assert masked_imaging_7x7.grids.pixelization.over_sample_size[0] == 4
+    assert masked_imaging_7x7.grids.over_sample_size_pixelization[0] == 4
 
 
 def test__grid_settings__sub_size(image_7x7, noise_map_7x7):
     dataset_7x7 = ds.AbstractDataset(
         data=image_7x7,
         noise_map=noise_map_7x7,
-        over_sampling=aa.OverSamplingDataset(
-            lp=2,
-            pixelization=4,
-        ),
+        over_sample_size_lp=2,
+        over_sample_size_pixelization=4,
     )
 
-    assert dataset_7x7.grids.lp.over_sample_size[0] == 2
-    assert dataset_7x7.grids.pixelization.over_sample_size[0] == 4
+    assert dataset_7x7.grids.over_sample_size_lp[0] == 2
+    assert dataset_7x7.grids.over_sample_size_pixelization[0] == 4
 
 
 def test__new_imaging_with_arrays_trimmed_via_kernel_shape():
@@ -140,10 +136,8 @@ def test__apply_over_sampling(image_7x7, noise_map_7x7):
     dataset_7x7 = aa.Imaging(
         data=image_7x7,
         noise_map=noise_map_7x7,
-        over_sampling=aa.OverSamplingDataset(
-            lp=2,
-            pixelization=2,
-        ),
+        over_sample_size_lp=2,
+        over_sample_size_pixelization=2,
     )
 
     # The grid and grid_pixelizaiton are a cached_property which needs to be reset,
@@ -159,14 +153,12 @@ def test__apply_over_sampling(image_7x7, noise_map_7x7):
     assert dataset_7x7.grids.pixelization[0][0] == pytest.approx(100.0, 1.0e-4)
 
     dataset_7x7 = dataset_7x7.apply_over_sampling(
-        over_sampling=aa.OverSamplingDataset(
-            lp=4,
-            pixelization=4,
-        )
+        over_sample_size_lp=2,
+        over_sample_size_pixelization=4,
     )
 
-    assert dataset_7x7.over_sampling.lp == 4
-    assert dataset_7x7.over_sampling.pixelization == 4
+    assert dataset_7x7.over_sample_size_lp.slim[0] == 2
+    assert dataset_7x7.over_sample_size_pixelization.slim[0] == 4
 
     assert dataset_7x7.grids.lp[0][0] == pytest.approx(3.0, 1.0e-4)
     assert dataset_7x7.grids.pixelization[0][0] == pytest.approx(3.0, 1.0e-4)
diff --git a/test_autoarray/dataset/plot/test_imaging_plotters.py b/test_autoarray/dataset/plot/test_imaging_plotters.py
index 6f71291f..d6c6673a 100644
--- a/test_autoarray/dataset/plot/test_imaging_plotters.py
+++ b/test_autoarray/dataset/plot/test_imaging_plotters.py
@@ -30,16 +30,16 @@ def test__individual_attributes_are_output(
         noise_map=True,
         psf=True,
         signal_to_noise_map=True,
-        over_sampling=True,
-        over_sampling_pixelization=True,
+        over_sample_size_lp=True,
+        over_sample_size_pixelization=True,
     )
 
     assert path.join(plot_path, "data.png") in plot_patch.paths
     assert path.join(plot_path, "noise_map.png") in plot_patch.paths
     assert path.join(plot_path, "psf.png") in plot_patch.paths
     assert path.join(plot_path, "signal_to_noise_map.png") in plot_patch.paths
-    assert path.join(plot_path, "over_sampling.png") in plot_patch.paths
-    assert path.join(plot_path, "over_sampling_pixelization.png") in plot_patch.paths
+    assert path.join(plot_path, "over_sample_size_lp.png") in plot_patch.paths
+    assert path.join(plot_path, "over_sample_size_pixelization.png") in plot_patch.paths
 
     plot_patch.paths = []
 
diff --git a/test_autoarray/inversion/plot/test_mapper_plotters.py b/test_autoarray/inversion/plot/test_mapper_plotters.py
index 00b72457..3dabe551 100644
--- a/test_autoarray/inversion/plot/test_mapper_plotters.py
+++ b/test_autoarray/inversion/plot/test_mapper_plotters.py
@@ -68,13 +68,14 @@ def test__get_2d__via_mapper_for_source_from(rectangular_mapper_7x7_3x3):
     assert mapper_plotter.visuals_2d.origin == None
     assert get_2d.origin.in_list == [(0.0, 0.0)]
     assert (
-        get_2d.grid
-        == rectangular_mapper_7x7_3x3.source_plane_data_grid.over_sampled
+        get_2d.grid == rectangular_mapper_7x7_3x3.source_plane_data_grid.over_sampled
     ).all()
     assert (get_2d.mesh_grid == rectangular_mapper_7x7_3x3.source_plane_mesh_grid).all()
-    border_grid = rectangular_mapper_7x7_3x3.mapper_grids.source_plane_data_grid.over_sampled[
-        rectangular_mapper_7x7_3x3.border_relocator.sub_border_slim
-    ]
+    border_grid = (
+        rectangular_mapper_7x7_3x3.mapper_grids.source_plane_data_grid.over_sampled[
+            rectangular_mapper_7x7_3x3.border_relocator.sub_border_slim
+        ]
+    )
     assert (get_2d.border == border_grid).all()
 
     include = aplt.Include2D(
diff --git a/test_autoarray/inversion/test_linear_obj.py b/test_autoarray/inversion/test_linear_obj.py
index 5267cd46..c5420435 100644
--- a/test_autoarray/inversion/test_linear_obj.py
+++ b/test_autoarray/inversion/test_linear_obj.py
@@ -44,9 +44,7 @@ def test__data_to_pix_unique_from():
 
     mask = aa.Mask2D.all_false(shape_native=(1, 2), pixel_scales=0.1)
 
-    grid = aa.Grid2D.uniform(
-        shape_native=(1, 2), pixel_scales=0.1, over_sample_size=2
-    )
+    grid = aa.Grid2D.uniform(shape_native=(1, 2), pixel_scales=0.1, over_sample_size=2)
 
     linear_obj = aa.AbstractLinearObjFuncList(grid=grid, regularization=None)
 
diff --git a/test_autoarray/plot/get_visuals/test_two_d.py b/test_autoarray/plot/get_visuals/test_two_d.py
index af576f41..5494bcf2 100644
--- a/test_autoarray/plot/get_visuals/test_two_d.py
+++ b/test_autoarray/plot/get_visuals/test_two_d.py
@@ -110,9 +110,11 @@ def test__via_mapper_for_source_from(rectangular_mapper_7x7_3x3):
         visuals_2d_via.grid
         == rectangular_mapper_7x7_3x3.source_plane_data_grid.over_sampled
     ).all()
-    border_grid = rectangular_mapper_7x7_3x3.mapper_grids.source_plane_data_grid.over_sampled[
-        rectangular_mapper_7x7_3x3.border_relocator.sub_border_slim
-    ]
+    border_grid = (
+        rectangular_mapper_7x7_3x3.mapper_grids.source_plane_data_grid.over_sampled[
+            rectangular_mapper_7x7_3x3.border_relocator.sub_border_slim
+        ]
+    )
     assert (visuals_2d_via.border == border_grid).all()
     assert (
         visuals_2d_via.mesh_grid == rectangular_mapper_7x7_3x3.source_plane_mesh_grid
diff --git a/test_autoarray/structures/grids/test_uniform_2d.py b/test_autoarray/structures/grids/test_uniform_2d.py
index 0fc83da4..4084908b 100644
--- a/test_autoarray/structures/grids/test_uniform_2d.py
+++ b/test_autoarray/structures/grids/test_uniform_2d.py
@@ -843,12 +843,7 @@ def test__is_uniform():
 
 
 def test__apply_over_sampling():
-
-    grid = aa.Grid2D.uniform(
-        shape_native=(2, 2),
-        pixel_scales=1.0,
-        over_sample_size=1
-    )
+    grid = aa.Grid2D.uniform(shape_native=(2, 2), pixel_scales=1.0, over_sample_size=1)
 
     grid = grid.apply_over_sampling(over_sample_size=2)