Improve docstrings for autoarray/fit package

autoarray/fit/fit_dataset.py

@@ -145,7 +145,7 @@ def __init__(
         noise_normalization
             The overall normalization term of the noise_map, summed over every data point.
         log_likelihood
-            The overall log likelihood of the model's fit to the dataset, summed over evey data point.
+            The overall log likelihood of the model's fit to the dataset, summed over every data point.
         """
         self.dataset = dataset
         self.use_mask_in_fit = use_mask_in_fit
@@ -162,6 +162,12 @@ def __init__(
 
     @property
     def _xp(self):
+        """
+        Returns the array module in use: `numpy` if JAX is disabled or `jax.numpy` if JAX is enabled.
+
+        This is controlled by the `use_jax` flag set during initialisation and is the single point of control
+        for switching between NumPy and JAX computation paths throughout the fit.
+        """
         if self.use_jax:
             import jax.numpy as jnp
 
@@ -170,10 +176,19 @@ def _xp(self):
 
     @property
     def mask(self) -> Mask2D:
+        """
+        The 2D mask of the dataset being fitted, where `False` entries are unmasked and included in the fit
+        and `True` entries are masked and excluded.
+        """
         return self.dataset.mask
 
     @property
     def grids(self) -> GridsInterface:
+        """
+        The grids of (y,x) coordinates associated with the dataset, adjusted by any `grid_offset` specified in
+        the `dataset_model`. Each grid (`lp`, `pixelization`, `blurring`) has the offset subtracted from it
+        before being returned.
+        """
 
         def subtracted_from(grid, offset):
             if grid is None:
@@ -200,10 +215,16 @@ def subtracted_from(grid, offset):
 
     @property
     def data(self) -> ty.DataLike:
+        """
+        The data of the dataset being fitted.
+        """
         return self.dataset.data
 
     @property
     def noise_map(self) -> ty.DataLike:
+        """
+        The noise-map of the dataset being fitted, representing the RMS noise in each pixel.
+        """
         return self.dataset.noise_map
 
     @property
@@ -310,19 +331,7 @@ def log_evidence(self) -> float:
         Log Evidence = -0.5*[Chi_Squared_Term + Regularization_Term + Log(Covariance_Regularization_Term) -
                            Log(Regularization_Matrix_Term) + Noise_Term]
 
-        Parameters
-        ----------
-        chi_squared
-            The chi-squared term of the inversion's fit to the data.
-        regularization_term
-            The regularization term of the inversion, which is the sum of the difference between reconstructed \
-            flux of every pixel multiplied by the regularization coefficient.
-        log_curvature_regularization_term
-            The log of the determinant of the sum of the curvature and regularization matrices.
-        log_regularization_term
-            The log of the determinant o the regularization matrix.
-        noise_normalization
-            The normalization noise_map-term for the data's noise-map.
+        Returns `None` if no inversion is present, in which case `log_likelihood` is used as the figure of merit.
         """
         if self.inversion is not None:
             return fit_util.log_evidence_from(
@@ -335,6 +344,11 @@ def log_evidence(self) -> float:
 
     @property
     def figure_of_merit(self) -> float:
+        """
+        The overall goodness-of-fit of the model to the dataset.
+
+        If the fit uses an inversion, this is the `log_evidence`; otherwise it is the `log_likelihood`.
+        """
         if self.inversion is not None:
             return self.log_evidence
 
@@ -371,4 +385,11 @@ def inversion(self) -> Optional[AbstractInversion]:
 
     @property
     def reduced_chi_squared(self) -> float:
+        """
+        The reduced chi-squared of the model's fit to the dataset, defined as:
+
+        Reduced_Chi_Squared = Chi_Squared / N_unmasked
+
+        where `N_unmasked` is the number of unmasked (i.e. `False`) pixels in the mask.
+        """
         return self.chi_squared / int(np.size(self.mask) - np.sum(self.mask))

autoarray/fit/fit_imaging.py

@@ -43,7 +43,7 @@ def __init__(
         noise_normalization
             The overall normalization term of the noise_map, summed over every data point.
         log_likelihood
-            The overall log likelihood of the model's fit to the dataset, summed over evey data point.
+            The overall log likelihood of the model's fit to the dataset, summed over every data point.
         """
 
         super().__init__(
@@ -55,8 +55,20 @@ def __init__(
 
     @property
     def data(self) -> ty.DataLike:
+        """
+        The imaging data being fitted, with any background sky level subtracted.
+
+        The background sky is taken from `dataset_model.background_sky_level`, which defaults to 0.0 if not
+        set, meaning this property is equivalent to `dataset.data` in the common case.
+        """
         return self.dataset.data - self.dataset_model.background_sky_level
 
     @property
     def blurred_image(self) -> Array2D:
+        """
+        The PSF-convolved (blurred) model image of the fit, as a 2D `Array2D`.
+
+        This is the model image after it has been convolved with the dataset's PSF. It must be implemented by
+        child classes (e.g. a tracer fit) that produce a blurred model image as part of their fitting procedure.
+        """
         raise NotImplementedError

autoarray/fit/fit_interferometer.py

@@ -24,15 +24,10 @@ def __init__(
 
         Parameters
         ----------
-        dataset : MaskedInterferometer
-            The masked interferometer dataset that is fitted.
+        dataset
+            The interferometer dataset that is fitted, containing the observed visibilities and noise-map.
         dataset_model
             Attributes which allow for parts of a dataset to be treated as a model (e.g. the background sky level).
-        model_data : Visibilities
-            The model visibilities the masked imaging is fitted with.
-        inversion : Inversion
-            If the fit uses an `Inversion` this is the instance of the object used to perform the fit. This determines
-            if the `log_likelihood` or `log_evidence` is used as the `figure_of_merit`.
         use_mask_in_fit
             If `True`, masked data points are omitted from the fit. If `False` they are not (in most use cases the
             `dataset` will have been processed to remove masked points, for example the `slim` representation).
@@ -51,7 +46,7 @@ def __init__(
         noise_normalization
             The overall normalization term of the noise_map, summed over every data point.
         log_likelihood
-            The overall log likelihood of the model's fit to the dataset, summed over evey data point.
+            The overall log likelihood of the model's fit to the dataset, summed over every data point.
         """
 
         super().__init__(
@@ -63,10 +58,22 @@ def __init__(
 
     @property
     def mask(self) -> np.ndarray:
+        """
+        The mask of the interferometer fit, returned as an all-`False` array matching the shape of the visibility data.
+
+        Interferometer data is not spatially masked in the same way as imaging data — all visibility measurements
+        are included in the fit — so this always returns an unmasked array.
+        """
         return np.full(shape=self.data.shape, fill_value=False)
 
     @property
     def transformer(self) -> ty.Transformer:
+        """
+        The Fourier transformer used to map between image space and visibility (uv-plane) space.
+
+        This is taken directly from the interferometer dataset and is used internally to compute the
+        `dirty_*` image-space representations of the fit quantities.
+        """
         return self.dataset.transformer
 
     @property
@@ -135,19 +142,9 @@ def log_evidence(self) -> float:
         Log Evidence = -0.5*[Chi_Squared_Term + Regularization_Term + Log(Covariance_Regularization_Term) -
                            Log(Regularization_Matrix_Term) + Noise_Term]
 
-        Parameters
-        ----------
-        chi_squared
-            The chi-squared term of the inversion's fit to the data.
-        regularization_term
-            The regularization term of the inversion, which is the sum of the difference between reconstructed \
-            flux of every pixel multiplied by the regularization coefficient.
-        log_curvature_regularization_term
-            The log of the determinant of the sum of the curvature and regularization matrices.
-        log_regularization_term
-            The log of the determinant o the regularization matrix.
-        noise_normalization
-            The normalization noise_map-term for the data's noise-map.
+        For interferometer fits the chi-squared uses the fast inversion chi-squared (`inversion.fast_chi_squared`).
+
+        Returns `None` if no inversion is present, in which case `log_likelihood` is used as the figure of merit.
         """
         if self.inversion is not None:
             return fit_util.log_evidence_from(
@@ -160,28 +157,56 @@ def log_evidence(self) -> float:
 
     @property
     def dirty_image(self) -> Array2D:
+        """
+        The dirty image of the observed visibility data, computed by applying the inverse Fourier transform to the
+        data visibilities. This is the image-space representation of the observed data before any deconvolution.
+        """
         return self.transformer.image_from(visibilities=self.data)
 
     @property
     def dirty_noise_map(self) -> Array2D:
+        """
+        The dirty noise-map, computed by applying the inverse Fourier transform to the noise-map visibilities.
+        This gives an image-space representation of the noise level across the field of view.
+        """
         return self.transformer.image_from(visibilities=self.noise_map)
 
     @property
     def dirty_signal_to_noise_map(self) -> Array2D:
+        """
+        The dirty signal-to-noise map, computed by applying the inverse Fourier transform to the signal-to-noise
+        visibilities. This gives an image-space representation of the signal-to-noise ratio across the field of view.
+        """
         return self.transformer.image_from(visibilities=self.signal_to_noise_map)
 
     @property
     def dirty_model_image(self) -> Array2D:
+        """
+        The dirty model image, computed by applying the inverse Fourier transform to the model data visibilities.
+        This is the image-space representation of the model before any deconvolution.
+        """
         return self.transformer.image_from(visibilities=self.model_data)
 
     @property
     def dirty_residual_map(self) -> Array2D:
+        """
+        The dirty residual map, computed by applying the inverse Fourier transform to the residual-map visibilities
+        (data - model_data). This is the image-space representation of the residuals.
+        """
         return self.transformer.image_from(visibilities=self.residual_map)
 
     @property
     def dirty_normalized_residual_map(self) -> Array2D:
+        """
+        The dirty normalized residual map, computed by applying the inverse Fourier transform to the
+        normalized residual-map visibilities ((data - model_data) / noise_map).
+        """
         return self.transformer.image_from(visibilities=self.normalized_residual_map)
 
     @property
     def dirty_chi_squared_map(self) -> Array2D:
+        """
+        The dirty chi-squared map, computed by applying the inverse Fourier transform to the chi-squared-map
+        visibilities (((data - model_data) / noise_map) ** 2.0).
+        """
         return self.transformer.image_from(visibilities=self.chi_squared_map)

autoarray/fit/fit_util.py

@@ -7,6 +7,14 @@
 
 
 def to_new_array(func):
+    """
+    Decorator that wraps the return value of a fit utility function in the same data structure as its first argument.
+
+    After computing the result, it calls `.with_new_array(result)` on the first keyword argument. If the first
+    argument does not have a `with_new_array` method (e.g. it is a plain ndarray), the raw result is returned
+    instead.
+    """
+
     @wraps(func)
     def wrapper(**kwargs):
         result = func(**kwargs)
@@ -28,8 +36,6 @@ def residual_map_from(*, data: ty.DataLike, model_data: ty.DataLike) -> ty.DataL
     ----------
     data
         The data that is fitted.
-    mask
-        The mask applied to the dataset, where `False` entries are included in the calculation.
     model_data
         The model data used to fit the data.
     """
@@ -50,8 +56,6 @@ def normalized_residual_map_from(
         The residual-map of the model-data fit to the dataset.
     noise_map
         The noise-map of the dataset.
-    mask
-        The mask applied to the residual-map, where `False` entries are included in the calculation.
     """
     return residual_map / noise_map
 
@@ -129,7 +133,7 @@ def chi_squared_map_complex_from(
     *, residual_map: np.ndarray, noise_map: np.ndarray
 ) -> np.ndarray:
     """
-    Returnss the chi-squared-map of the fit of complex model-data to a dataset, where:
+    Returns the chi-squared-map of the fit of complex model-data to a dataset, where:
 
     Chi_Squared = ((Residuals) / (Noise)) ** 2.0 = ((Data - Model)**2.0)/(Variances)
 
@@ -229,7 +233,7 @@ def chi_squared_map_with_mask_from(
     *, residual_map: ty.DataLike, noise_map: ty.DataLike, mask: Mask, xp=np
 ) -> ty.DataLike:
     """
-    Returnss the chi-squared-map of the fit of model-data to a masked dataset, where:
+    Returns the chi-squared-map of the fit of model-data to a masked dataset, where:
 
     Chi_Squared = ((Residuals) / (Noise)) ** 2.0 = ((Data - Model)**2.0)/(Variances)
 
@@ -289,10 +293,14 @@ def chi_squared_with_mask_fast_from(
 
     Parameters
     ----------
-    chi_squared_map
-        The chi-squared-map of values of the model-data fit to the dataset.
+    data
+        The data that is fitted.
     mask
-        The mask applied to the chi-squared-map, where `False` entries are included in the calculation.
+        The mask applied to the dataset, where `False` entries are included in the calculation.
+    model_data
+        The model data used to fit the data.
+    noise_map
+        The noise-map of the dataset.
     """
     return float(
         xp.sum(
@@ -412,7 +420,7 @@ def log_evidence_from(
     log_curvature_regularization_term
         The log of the determinant of the sum of the curvature and regularization matrices.
     log_regularization_term
-        The log of the determinant o the regularization matrix.
+        The log of the determinant of the regularization matrix.
     noise_normalization
         The normalization noise_map-term for the dataset's noise-map.
     """
@@ -447,7 +455,7 @@ def residual_flux_fraction_map_with_mask_from(
     *, residual_map: np.ndarray, data: np.ndarray, mask: Mask, xp=np
 ) -> np.ndarray:
     """
-    Returnss the residual flux fraction map of the fit of model-data to a masked dataset, where:
+    Returns the residual flux fraction map of the fit of model-data to a masked dataset, where:
 
     Residual_Flux_Fraction = Residuals / Data = (Data - Model)/Data