docs: refactor and complete docstrings for autoarray/dataset

autoarray/dataset/abstract/dataset.py

@@ -134,34 +134,62 @@ def __init__(
 
     @property
     def grid(self):
+        """
+        The primary coordinate grid of the dataset, equivalent to `grids.lp`.
+
+        Returns the light-profile `Grid2D` aligned with the centres of all unmasked image pixels.
+        This is the grid used for the majority of model calculations (e.g. evaluating galaxy light
+        profiles).
+        """
         return self.grids.lp
 
     @property
     def shape_native(self):
+        """
+        The 2D shape of the dataset image in its native (unmasked) dimensions, e.g. (rows, columns).
+        """
         return self.mask.shape_native
 
     @property
     def shape_slim(self):
+        """
+        The 1D size of the dataset data array after masking, i.e. the number of unmasked pixels.
+        """
         return self.data.shape_slim
 
     @property
     def pixel_scales(self):
+        """
+        The (y, x) arcsecond-to-pixel conversion factor of the dataset, as a (float, float) tuple.
+        """
         return self.mask.pixel_scales
 
     @property
     def mask(self) -> Union[Mask1D, Mask2D]:
+        """
+        The mask of the dataset, derived from the mask of the `data` array.
+        """
         return self.data.mask
 
     def apply_over_sampling(self):
+        """
+        Apply new over-sampling sizes to the dataset grids.
+
+        Subclasses must implement this method to rebuild the `GridsDataset` with updated
+        `over_sample_size_lp` and `over_sample_size_pixelization` values.
+        """
         raise NotImplementedError
 
     @property
     def signal_to_noise_map(self) -> Structure:
         """
-        The estimated signal-to-noise_maps mappers of the image.
+        The signal-to-noise map of the dataset, computed as `data / noise_map`.
 
-        Warnings arise when masked native noise-maps are used, whose masked entries are given values of 0.0. We
-        use the warnings module to suppress these RunTimeWarnings.
+        Values below zero are clamped to zero, as negative signal-to-noise is not physically
+        meaningful (it indicates the data is below zero due to noise, not a real negative signal).
+
+        RuntimeWarnings from dividing by zero in masked pixels (where the noise map is 0.0) are
+        suppressed, as these masked values are never used in downstream calculations.
         """
         warnings.filterwarnings("ignore")
 
@@ -172,7 +200,7 @@ def signal_to_noise_map(self) -> Structure:
     @property
     def signal_to_noise_max(self) -> float:
         """
-        The maximum value of signal-to-noise_maps in an image pixel in the image's signal-to-noise_maps mappers.
+        The maximum signal-to-noise value across all unmasked pixels in the dataset.
         """
         return np.max(self.signal_to_noise_map)
 
@@ -185,6 +213,27 @@ def noise_covariance_matrix_inv(self) -> np.ndarray:
         return np.linalg.inv(self.noise_covariance_matrix)
 
     def trimmed_after_convolution_from(self, kernel_shape) -> "AbstractDataset":
+        """
+        Return a copy of the dataset with all arrays trimmed to remove the border pixels affected
+        by PSF convolution edge effects.
+
+        When a model image is convolved with a PSF kernel, the pixels at the border of the image
+        cannot be correctly convolved because they lack sufficient neighbouring pixels. These border
+        pixels have unreliable values after convolution. This method trims the `data`, `noise_map`,
+        `over_sample_size_lp` and `over_sample_size_pixelization` arrays by the kernel half-width
+        on each side, so that only pixels with a complete convolution kernel neighbourhood remain.
+
+        Parameters
+        ----------
+        kernel_shape
+            The (rows, cols) shape of the PSF convolution kernel. The dataset arrays are trimmed
+            by `kernel_shape // 2` pixels on each side in each dimension.
+
+        Returns
+        -------
+        AbstractDataset
+            A shallow copy of the dataset with all arrays trimmed to the post-convolution shape.
+        """
         dataset = copy.copy(self)
 
         dataset.data = dataset.data.trimmed_after_convolution_from(

autoarray/dataset/grids.py

@@ -44,6 +44,9 @@ def __init__(
 
         Parameters
         ----------
+        mask
+            The 2D mask defining which pixels are included in the dataset. All grids are constructed
+            to align with the centres of the unmasked pixels in this mask.
         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
@@ -53,7 +56,8 @@ def __init__(
             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.
+            via 2D convolution. Required to compute the blurring grid; if `None` the blurring grid
+            is not constructed.
         """
         self.mask = mask
         self.over_sample_size_lp = over_sample_size_lp
@@ -67,7 +71,17 @@ def __init__(
 
     @property
     def lp(self):
+        """
+        The light-profile grid: a `Grid2D` of (y,x) Cartesian coordinates at the centre of every
+        unmasked image pixel, used for evaluating light profiles and other spatial calculations
+        during model fitting.
+
+        The grid uses `over_sample_size_lp` to perform over-sampled sub-pixel integration,
+        approximating the 2D line integral of the light profile within each pixel. This grid is
+        what most model-fitting calculations use (e.g. computing galaxy images).
 
+        This property is lazily evaluated and cached on first access.
+        """
         if self._lp is not None:
             return self._lp
 
@@ -80,6 +94,16 @@ def lp(self):
 
     @property
     def pixelization(self):
+        """
+        The pixelization grid: a `Grid2D` of (y,x) Cartesian coordinates at the centre of every
+        unmasked image pixel, dedicated to pixelized source reconstructions via the `inversion` module.
+
+        This grid uses `over_sample_size_pixelization` which can differ from `over_sample_size_lp`,
+        allowing the pixelization to benefit from a different (e.g. lower) over-sampling resolution
+        than the light-profile grid.
+
+        This property is lazily evaluated and cached on first access.
+        """
         if self._pixelization is not None:
             return self._pixelization
 
@@ -92,7 +116,18 @@ def pixelization(self):
 
     @property
     def blurring(self):
+        """
+        The blurring grid: a `Grid2D` of (y,x) coordinates for pixels that lie just outside the
+        mask but whose light can be scattered into the unmasked region by the PSF.
+
+        When convolving a model image with the PSF, pixels neighbouring the mask boundary can
+        contribute flux to unmasked pixels. The blurring grid provides the coordinates of these
+        border pixels so their light profile values can be evaluated and included in the convolution.
+
+        Returns `None` if no PSF was supplied (i.e. no blurring is performed).
 
+        This property is lazily evaluated and cached on first access.
+        """
         if self._blurring is not None:
             return self._blurring
 
@@ -113,7 +148,16 @@ def blurring(self):
 
     @property
     def border_relocator(self) -> BorderRelocator:
+        """
+        The border relocator for the pixelization grid.
+
+        During pixelized source reconstruction, source-plane coordinates that map outside the
+        border of the pixelization mesh can cause numerical problems. The `BorderRelocator`
+        detects these coordinates and relocates them to the border of the mesh, preventing
+        ill-conditioned inversions.
 
+        This property is lazily evaluated and cached on first access.
+        """
         if self._border_relocator is not None:
             return self._border_relocator
 
@@ -133,6 +177,26 @@ def __init__(
         blurring=None,
         border_relocator=None,
     ):
+        """
+        A lightweight plain-data container for pre-constructed dataset grids.
+
+        Unlike `GridsDataset`, this class performs no computation — it simply holds grids that have
+        already been created elsewhere. It is used in test fixtures and mock datasets where a full
+        `GridsDataset` is not needed, but code that accesses `dataset.grids.lp` or
+        `dataset.grids.pixelization` still needs to work.
+
+        Parameters
+        ----------
+        lp
+            The light-profile `Grid2D` used for evaluating light profiles during model fitting.
+        pixelization
+            The pixelization `Grid2D` used for source reconstruction via the inversion module.
+        blurring
+            The blurring `Grid2D` for pixels outside the mask that contribute flux via PSF convolution.
+        border_relocator
+            The `BorderRelocator` used to remap out-of-bounds source-plane coordinates to the
+            pixelization mesh border.
+        """
         self.lp = lp
         self.pixelization = pixelization
         self.blurring = blurring

autoarray/dataset/imaging/dataset.py

@@ -68,6 +68,10 @@ def __init__(
         psf
             The Point Spread Function kernel of the image which accounts for diffraction due to the telescope optics
             via 2D convolution.
+        psf_setup_state
+            If `True`, a `ConvolverState` is precomputed from the PSF kernel and mask, storing the
+            convolution pair indices required for efficient 2D convolution. This is set automatically
+            to `True` when a mask is applied via `apply_mask()` and should not normally be set by hand.
         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.
@@ -236,13 +240,25 @@ def apply_mask(self, mask: Mask2D) -> "Imaging":
         Apply a mask to the imaging dataset, whereby the mask is applied to the image data, noise-map and other
         quantities one-by-one.
 
-        The `apply_mask` function cannot be called multiple times, if it is a mask may remove data, therefore
-        an exception is raised. If you wish to apply a new mask, reload the dataset from .fits files.
+        The mask is applied to the `data`, `noise_map`, `over_sample_size_lp` and
+        `over_sample_size_pixelization` arrays. If a `noise_covariance_matrix` is present, the rows
+        and columns corresponding to masked pixels are removed so it stays consistent with the
+        remaining unmasked pixels. The PSF `ConvolverState` is recomputed for the new mask.
+
+        The `apply_mask` function cannot be called multiple times — a new mask cannot expand the
+        unmasked region beyond what was already unmasked, as the underlying data has already been
+        trimmed. An exception is raised if this is attempted. If you wish to apply a different mask,
+        reload the dataset from .fits files.
 
         Parameters
         ----------
         mask
             The 2D mask that is applied to the image.
+
+        Returns
+        -------
+        Imaging
+            A new `Imaging` dataset with the mask applied to all arrays.
         """
         invalid = np.logical_and(self.data.mask, np.logical_not(mask))
 
@@ -413,22 +429,27 @@ def apply_sparse_operator(
         batch_size: int = 128,
     ):
         """
+        Precompute the PSF precision operator for efficient pixelized source reconstruction.
+
         The sparse linear algebra formalism precomputes the convolution of every pair of masked
-        noise-map values given the PSF (see `inversion.inversion_util`).
+        noise-map values given the PSF (see `inversion.inversion_util`). This is the imaging
+        equivalent of the interferometer NUFFT precision matrix.
 
-        The `WTilde` object stores these precomputed values in the imaging dataset ensuring they are only computed once
-        per analysis.
+        The `ImagingSparseOperator` stores these precomputed values in the imaging dataset ensuring
+        they are only computed once per analysis, enabling fast repeated likelihood evaluations during
+        model fitting.
 
-        This uses lazy allocation such that the calculation is only performed when the wtilde matrices are used,
-        ensuring efficient set up of the `Imaging` class.
+        Parameters
+        ----------
+        batch_size
+            The number of image pixels processed per batch when computing the sparse operator via
+            FFT-based convolution. Reducing this lowers peak memory usage at the cost of speed.
 
         Returns
         -------
-        batch_size
-            The size of batches used to compute the w-tilde curvature matrix via FFT-based convolution,
-            which can be reduced to produce lower memory usage at the cost of speed
-        use_jax
-            Whether to use JAX to compute W-Tilde. This requires JAX to be installed.
+        Imaging
+            A new `Imaging` dataset with the precomputed `ImagingSparseOperator` attached, enabling
+            efficient pixelized source reconstruction via the sparse linear algebra formalism.
         """
 
         logger.info(
@@ -459,22 +480,20 @@ def apply_sparse_operator_cpu(
         self,
     ):
         """
-        The sparse linear algebra formalism precomputes the convolution of every pair of masked
-        noise-map values given the PSF (see `inversion.inversion_util`).
+        Precompute the PSF precision operator using a CPU-only Numba implementation.
 
-        The `WTilde` object stores these precomputed values in the imaging dataset ensuring they are only computed once
-        per analysis.
+        This is the CPU alternative to `apply_sparse_operator()`, using Numba JIT compilation
+        for the convolution loop rather than JAX. It requires `numba` to be installed; an
+        `InversionException` is raised if it is not available.
 
-        This uses lazy allocation such that the calculation is only performed when the wtilde matrices are used,
-        ensuring efficient set up of the `Imaging` class.
+        The resulting `SparseLinAlgImagingNumba` operator is stored on the returned `Imaging`
+        dataset and used by `FitImaging` when performing pixelized source reconstructions.
 
         Returns
         -------
-        batch_size
-            The size of batches used to compute the w-tilde curvature matrix via FFT-based convolution,
-            which can be reduced to produce lower memory usage at the cost of speed.
-        use_jax
-            Whether to use JAX to compute W-Tilde. This requires JAX to be installed.
+        Imaging
+            A new `Imaging` dataset with a precomputed Numba-based sparse operator attached,
+            enabling efficient pixelized source reconstruction on CPU hardware.
         """
         try:
             import numba

autoarray/dataset/imaging/simulator.py

@@ -33,8 +33,8 @@ def __init__(
 
         The simulation of an `Imaging` dataset uses the following steps:
 
-        1) Receive as input a raw image of what the data looks like before any simulaiton process is applied.
-        2) Include dirrection due to the telescope optics by convolve the image with an input Point Spread
+        1) Receive as input a raw image of what the data looks like before any simulation process is applied.
+        2) Include diffraction due to the telescope optics by convolving the image with an input Point Spread
            Function (PSF).
         3) Use input values of the background sky level in every pixel of the image to add the background sky to
            the PSF convolved image.
@@ -121,7 +121,18 @@ def via_image_from(
         Parameters
         ----------
         image
-            The 2D image from which the Imaging dataset is simulated.
+            The 2D image from which the Imaging dataset is simulated (e.g. a model galaxy image
+            before any telescope effects are applied). Must be an `Array2D`.
+        over_sample_size
+            If provided, the returned dataset has its over-sampling updated via `apply_over_sampling`.
+            Should be an `Array2D` of integer sub-grid sizes with the same shape as the image.
+        xp
+            The array module to use for PSF convolution (default `np` for NumPy, or `jnp` for JAX).
+
+        Returns
+        -------
+        Imaging
+            The simulated imaging dataset with PSF convolution, noise and background sky applied.
         """
 
         exposure_time_map = Array2D.full(