Interferometer Inversion implemented and added to tracer / fit modules

Author: Jammy2211

autoarray/operators/inversion/inversions.py

@@ -1,21 +1,110 @@
 import numpy as np
 
 from autoarray import exc
+from autoarray.structures import visibilities as vis
+from autoarray.masked import masked_dataset as md
 from autoarray.util import inversion_util
 
 
 def inversion(masked_dataset, mapper, regularization):
 
-    return InversionImaging.from_data_mapper_and_regularization(
-        image=masked_dataset.image,
-        noise_map=masked_dataset.noise_map,
-        convolver=masked_dataset.convolver,
-        mapper=mapper,
-        regularization=regularization,
-    )
+    if isinstance(masked_dataset, md.MaskedImaging):
 
+        return InversionImaging.from_data_mapper_and_regularization(
+            image=masked_dataset.image,
+            noise_map=masked_dataset.noise_map,
+            convolver=masked_dataset.convolver,
+            mapper=mapper,
+            regularization=regularization,
+        )
+    
+    elif isinstance(masked_dataset, md.MaskedInterferometer):
+
+        return InversionInterferometer.from_data_mapper_and_regularization(
+            visibilities=masked_dataset.visibilities,
+            noise_map=masked_dataset.noise_map,
+            transformer=masked_dataset.transformer,
+            mapper=mapper,
+            regularization=regularization,
+        )
+
+
+class Inversion(object):
+
+    def __init__(
+        self,
+        noise_map,
+        mapper,
+        regularization,
+        regularization_matrix,
+        curvature_reg_matrix,
+        reconstruction,
+    ):
+
+        self.noise_map = noise_map
+        self.mapper = mapper
+        self.regularization = regularization
+        self.regularization_matrix = regularization_matrix
+        self.curvature_reg_matrix = curvature_reg_matrix
+        self.reconstruction = reconstruction
+
+    @property
+    def errors_with_covariance(self):
+        return np.linalg.inv(self.curvature_reg_matrix)
+
+    @property
+    def errors(self):
+        return np.diagonal(self.errors_with_covariance)
+
+    @property
+    def regularization_term(self):
+        """ Compute the regularization term of an inversion. This term represents the sum of the difference in flux \
+        between every pair of neighboring pixels. This is computed as:
+
+        s_T * H * s = solution_vector.T * regularization_matrix * solution_vector
+
+        The term is referred to as *G_l* in Warren & Dye 2003, Nightingale & Dye 2015.
+
+        The above works include the regularization_matrix coefficient (lambda) in this calculation. In PyAutoLens, \
+        this is already in the regularization matrix and thus implicitly included in the matrix multiplication.
+        """
+        return np.matmul(
+            self.reconstruction.T,
+            np.matmul(self.regularization_matrix, self.reconstruction),
+        )
 
-class InversionImaging(object):
+    @property
+    def log_det_curvature_reg_matrix_term(self):
+        return self.log_determinant_of_matrix_cholesky(self.curvature_reg_matrix)
+
+    @property
+    def log_det_regularization_matrix_term(self):
+        return self.log_determinant_of_matrix_cholesky(self.regularization_matrix)
+
+    @staticmethod
+    def log_determinant_of_matrix_cholesky(matrix):
+        """There are two terms in the inversion's Bayesian likelihood function which require the log determinant of \
+        a matrix. These are (Nightingale & Dye 2015, Nightingale, Dye and Massey 2018):
+
+        ln[det(F + H)] = ln[det(curvature_reg_matrix)]
+        ln[det(H)]     = ln[det(regularization_matrix)]
+
+        The curvature_reg_matrix is positive-definite, which means the above log determinants can be computed \
+        efficiently (compared to using np.det) by using a Cholesky decomposition first and summing the log of each \
+        diagonal term.
+
+        Parameters
+        -----------
+        matrix : ndarray
+            The positive-definite matrix the log determinant is computed for.
+        """
+        try:
+            return 2.0 * np.sum(np.log(np.diag(np.linalg.cholesky(matrix))))
+        except np.linalg.LinAlgError:
+            raise exc.InversionException()
+
+
+class InversionImaging(Inversion):
     def __init__(
         self,
         image,
@@ -63,14 +152,12 @@ def __init__(
             The vector containing the reconstructed fit to the hyper_galaxies.
         """
 
+        super(InversionImaging, self).__init__(noise_map=noise_map, mapper=mapper, regularization=regularization,
+                                               regularization_matrix=regularization_matrix, curvature_reg_matrix=curvature_reg_matrix,
+                                               reconstruction=reconstruction)
+
         self.image = image
-        self.noise_map = noise_map
-        self.mapper = mapper
-        self.regularization = regularization
         self.blurred_mapping_matrix = blurred_mapping_matrix
-        self.regularization_matrix = regularization_matrix
-        self.curvature_reg_matrix = curvature_reg_matrix
-        self.reconstruction = reconstruction
 
     @classmethod
     def from_data_mapper_and_regularization(
@@ -115,22 +202,14 @@ def from_data_mapper_and_regularization(
 
     @property
     def mapped_reconstructed_image(self):
-        reconstructed_image = inversion_util.mapper_reconstructed_image_from_blurred_mapping_matrix_and_reconstruction(
-            blurred_mapping_matrix=self.blurred_mapping_matrix,
+        reconstructed_image = inversion_util.mapped_reconstructed_data_from_mapping_matrix_and_reconstruction(
+            mapping_matrix=self.blurred_mapping_matrix,
             reconstruction=self.reconstruction,
         )
         return self.mapper.grid.mapping.array_from_array_1d(
             array_1d=reconstructed_image
         )
 
-    @property
-    def errors_with_covariance(self):
-        return np.linalg.inv(self.curvature_reg_matrix)
-
-    @property
-    def errors(self):
-        return np.diagonal(self.errors_with_covariance)
-
     @property
     def residual_map(self):
         return inversion_util.inversion_residual_map_from_pixelization_values_and_reconstructed_data_1d(
@@ -160,49 +239,129 @@ def chi_squared_map(self):
             all_sub_mask_1d_indexes_for_pixelization_1d_index=self.mapper.all_sub_mask_1d_indexes_for_pixelization_1d_index,
         )
 
-    @property
-    def regularization_term(self):
-        """ Compute the regularization term of an inversion. This term represents the sum of the difference in flux \
-        between every pair of neighboring pixels. This is computed as:
 
-        s_T * H * s = solution_vector.T * regularization_matrix * solution_vector
+class InversionInterferometer(Inversion):
+    def __init__(
+        self,
+        visibilities,
+        noise_map,
+        mapper,
+        regularization,
+        transformed_mapping_matrices,
+        regularization_matrix,
+        curvature_reg_matrix,
+        reconstruction,
+    ):
+        """ An inversion, which given an input image and noise-map reconstructs the image using a linear inversion, \
+        including a convolution that accounts for blurring.
 
-        The term is referred to as *G_l* in Warren & Dye 2003, Nightingale & Dye 2015.
+        The inversion uses a 2D pixelization to perform the reconstruction by util each pixelization pixel to a \
+        set of image pixels via a mapper. The reconstructed pixelization is smoothed via a regularization scheme to \
+        prevent over-fitting noise.
 
-        The above works include the regularization_matrix coefficient (lambda) in this calculation. In PyAutoLens, \
-        this is already in the regularization matrix and thus implicitly included in the matrix multiplication.
+        Parameters
+        -----------
+        image_1d : ndarray
+            Flattened 1D array of the observed image the inversion is fitting.
+        noise_map : ndarray
+            Flattened 1D array of the noise-map used by the inversion during the fit.   
+        convolver : imaging.convolution.Convolver
+            The convolver used to blur the util matrix with the PSF.
+        mapper : inversion.mappers.Mapper
+            The util between the image-pixels (via its / sub-grid) and pixelization pixels.
+        regularization : inversion.regularization.Regularization
+            The regularization scheme applied to smooth the pixelization used to reconstruct the image for the \
+            inversion
+
+        Attributes
+        -----------
+        blurred_mapping_matrix : ndarray
+            The matrix representing the blurred mappings between the image's sub-grid of pixels and the pixelization \
+            pixels.
+        regularization_matrix : ndarray
+            The matrix defining how the pixelization's pixels are regularized with one another for smoothing (H).
+        curvature_matrix : ndarray
+            The curvature_matrix between each pixelization pixel and all other pixelization pixels (F).
+        curvature_reg_matrix : ndarray
+            The curvature_matrix + regularization matrix.
+        solution_vector : ndarray
+            The vector containing the reconstructed fit to the hyper_galaxies.
         """
-        return np.matmul(
-            self.reconstruction.T,
-            np.matmul(self.regularization_matrix, self.reconstruction),
+
+        super(InversionInterferometer, self).__init__(noise_map=noise_map, mapper=mapper, regularization=regularization,
+                                               regularization_matrix=regularization_matrix, curvature_reg_matrix=curvature_reg_matrix,
+                                               reconstruction=reconstruction)
+
+        self.visibilities = visibilities
+        self.transformed_mapping_matrices = transformed_mapping_matrices
+
+    @classmethod
+    def from_data_mapper_and_regularization(
+        cls, visibilities, noise_map, transformer, mapper, regularization
+    ):
+
+        transformed_mapping_matrices = transformer.transformed_mapping_matrices_from_mapping_matrix(
+            mapping_matrix=mapper.mapping_matrix
         )
 
-    @property
-    def log_det_curvature_reg_matrix_term(self):
-        return self.log_determinant_of_matrix_cholesky(self.curvature_reg_matrix)
+        real_data_vector = inversion_util.data_vector_from_transformed_mapping_matrix_and_data(
+            transformed_mapping_matrix=transformed_mapping_matrices[0],
+            visibilities=visibilities[:,0],
+            noise_map=noise_map[:,0],
+        )
 
-    @property
-    def log_det_regularization_matrix_term(self):
-        return self.log_determinant_of_matrix_cholesky(self.regularization_matrix)
+        imag_data_vector = inversion_util.data_vector_from_transformed_mapping_matrix_and_data(
+            transformed_mapping_matrix=transformed_mapping_matrices[1],
+            visibilities=visibilities[:,1],
+            noise_map=noise_map[:,1],
+        )
 
-    @staticmethod
-    def log_determinant_of_matrix_cholesky(matrix):
-        """There are two terms in the inversion's Bayesian likelihood function which require the log determinant of \
-        a matrix. These are (Nightingale & Dye 2015, Nightingale, Dye and Massey 2018):
+        real_curvature_matrix = inversion_util.curvature_matrix_from_transformed_mapping_matrix(
+            transformed_mapping_matrix=transformed_mapping_matrices[0],
+            noise_map=noise_map[:,0],
+        )
 
-        ln[det(F + H)] = ln[det(curvature_reg_matrix)]
-        ln[det(H)]     = ln[det(regularization_matrix)]
+        imag_curvature_matrix = inversion_util.curvature_matrix_from_transformed_mapping_matrix(
+            transformed_mapping_matrix=transformed_mapping_matrices[1],
+            noise_map=noise_map[:,1],
+        )
 
-        The curvature_reg_matrix is positive-definite, which means the above log determinants can be computed \
-        efficiently (compared to using np.det) by using a Cholesky decomposition first and summing the log of each \
-        diagonal term.
+        regularization_matrix = regularization.regularization_matrix_from_mapper(
+            mapper=mapper
+        )
+
+        real_curvature_reg_matrix = np.add(real_curvature_matrix, regularization_matrix)
+        imag_curvature_reg_matrix = np.add(imag_curvature_matrix, regularization_matrix)
+
+        data_vector = np.add(real_data_vector, imag_data_vector)
+        curvature_reg_matrix = np.add(real_curvature_reg_matrix, imag_curvature_reg_matrix)
 
-        Parameters
-        -----------
-        matrix : ndarray
-            The positive-definite matrix the log determinant is computed for.
-        """
         try:
-            return 2.0 * np.sum(np.log(np.diag(np.linalg.cholesky(matrix))))
+            values = np.linalg.solve(curvature_reg_matrix, data_vector)
         except np.linalg.LinAlgError:
             raise exc.InversionException()
+
+        return InversionInterferometer(
+            visibilities=visibilities,
+            noise_map=noise_map,
+            mapper=mapper,
+            regularization=regularization,
+            transformed_mapping_matrices=transformed_mapping_matrices,
+            regularization_matrix=regularization_matrix,
+            curvature_reg_matrix=curvature_reg_matrix,
+            reconstruction=values,
+        )
+
+    @property
+    def mapped_reconstructed_visibilities(self):
+        real_visibilities = inversion_util.mapped_reconstructed_data_from_mapping_matrix_and_reconstruction(
+            mapping_matrix=self.transformed_mapping_matrices[0],
+            reconstruction=self.reconstruction,
+        )
+
+        imag_visibilities = inversion_util.mapped_reconstructed_data_from_mapping_matrix_and_reconstruction(
+            mapping_matrix=self.transformed_mapping_matrices[1],
+            reconstruction=self.reconstruction,
+        )
+
+        return vis.Visibilities(visibilities_1d=np.stack((real_visibilities, imag_visibilities), axis=-1))

autoarray/util/inversion_util.py

@@ -101,25 +101,25 @@ def curvature_matrix_from_blurred_mapping_matrix_jit(
 
 
 @decorator_util.jit()
-def mapper_reconstructed_image_from_blurred_mapping_matrix_and_reconstruction(
-    blurred_mapping_matrix, reconstruction
+def mapped_reconstructed_data_from_mapping_matrix_and_reconstruction(
+    mapping_matrix, reconstruction
 ):
     """ Compute the reconstructed hyper_galaxies vector from the blurrred util matrix *f* and solution vector *S*.
 
     Parameters
     -----------
-    blurred_mapping_matrix : ndarray
+    mapping_matrix : ndarray
         The matrix representing the blurred mappings between sub-grid pixels and pixelization pixels.
 
     """
-    reconstructed_data_vector = np.zeros(blurred_mapping_matrix.shape[0])
-    for i in range(blurred_mapping_matrix.shape[0]):
+    mapped_reconstructred_data = np.zeros(mapping_matrix.shape[0])
+    for i in range(mapping_matrix.shape[0]):
         for j in range(reconstruction.shape[0]):
-            reconstructed_data_vector[i] += (
-                    reconstruction[j] * blurred_mapping_matrix[i, j]
+            mapped_reconstructred_data[i] += (
+                    reconstruction[j] * mapping_matrix[i, j]
             )
 
-    return reconstructed_data_vector
+    return mapped_reconstructred_data
 
 
 @decorator_util.jit()
@@ -142,12 +142,10 @@ def data_vector_from_transformed_mapping_matrix_and_data(
     data_vector = np.zeros(transformed_mapping_matrix.shape[1])
 
     for vis_1d_index in range(transformed_mapping_matrix.shape[0]):
-        for pix_1_index in range(transformed_mapping_matrix.shape[1]):
-            data_vector[pix_1_index] += (
-                    visibilities[vis_1d_index]
-                    * transformed_mapping_matrix[vis_1d_index, pix_1_index]
+        for pix_1d_index in range(transformed_mapping_matrix.shape[1]):
+            data_vector[pix_1d_index] += visibilities[vis_1d_index] \
+                    * transformed_mapping_matrix[vis_1d_index, pix_1d_index] \
                     / (noise_map[vis_1d_index] ** 2.0)
-            )
 
     return data_vector