From 21a1744cd8ef00d906769251936772ef2f480cb8 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Mon, 9 Mar 2026 16:22:40 +0000 Subject: [PATCH] improve docstrings for autoarray/fit package MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - fit_util.py: add docstring to `to_new_array` decorator; remove spurious `mask` param from `residual_map_from` and `normalized_residual_map_from`; fix wrong Parameters in `chi_squared_with_mask_fast_from` (was listing `chi_squared_map`, now lists the actual params: data/mask/model_data/noise_map); fix three "Returnss" typos; fix "o the" typo in `log_evidence_from` - fit_dataset.py: add docstrings to `_xp`, `mask`, `grids`, `data`, `noise_map`, `figure_of_merit`, `reduced_chi_squared`; remove wrong Parameters section from `log_evidence` property; fix "evey" typo - fit_imaging.py: add docstrings to `data` and `blurred_image`; fix "evey" typo - fit_interferometer.py: fix `__init__` parameters (remove non-parameter items `model_data`/`inversion`, fix `MaskedInterferometer` → `Interferometer`); add docstrings to `mask`, `transformer`, and all seven `dirty_*` properties; fix `log_evidence` wrong Parameters section; fix "evey" typo Co-Authored-By: Claude Sonnet 4.6 --- autoarray/fit/fit_dataset.py | 49 +++++++++++++++------ autoarray/fit/fit_imaging.py | 14 +++++- autoarray/fit/fit_interferometer.py | 67 ++++++++++++++++++++--------- autoarray/fit/fit_util.py | 30 ++++++++----- 4 files changed, 113 insertions(+), 47 deletions(-) diff --git a/autoarray/fit/fit_dataset.py b/autoarray/fit/fit_dataset.py index d97e736e5..8ef8dfc64 100644 --- a/autoarray/fit/fit_dataset.py +++ b/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)) diff --git a/autoarray/fit/fit_imaging.py b/autoarray/fit/fit_imaging.py index bbf22c3d9..57f05af82 100644 --- a/autoarray/fit/fit_imaging.py +++ b/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 diff --git a/autoarray/fit/fit_interferometer.py b/autoarray/fit/fit_interferometer.py index 8ef3451c7..e849b07c0 100644 --- a/autoarray/fit/fit_interferometer.py +++ b/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) diff --git a/autoarray/fit/fit_util.py b/autoarray/fit/fit_util.py index 1b126a2d2..01e9a621e 100644 --- a/autoarray/fit/fit_util.py +++ b/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