from functools import partial
from jax import jit, lax, numpy as jnp
import jax
import jaxtronomy.ImSim.de_lens as de_lens
from jaxtronomy.ImSim.image_model import ImageModel
__all__ = ["ImageLinearFit"]
[docs]
class ImageLinearFit(ImageModel):
"""Linear version class, inherits ImageModel.
When light models use pixel-based profile types, such as 'SLIT_STARLETS', the WLS
linear inversion is replaced by the regularized inversion performed by an external
solver. The current pixel-based solver is provided by the SLITronomy plug-in.
"""
[docs]
def __init__(
self,
data_class,
psf_class,
lens_model_class=None,
source_model_class=None,
lens_light_model_class=None,
point_source_class=None,
extinction_class=None,
kwargs_numerics=None,
likelihood_mask=None,
psf_error_map_bool_list=None,
kwargs_pixelbased=None,
):
"""
:param data_class: ImageData() instance
:param psf_class: PSF() instance
:param lens_model_class: LensModel() instance
:param source_model_class: LightModel() instance
:param lens_light_model_class: LightModel() instance
:param point_source_class: PointSource() instance
:param extinction_class: DifferentialExtinction() instance
:param kwargs_numerics: keyword arguments passed to the Numerics module
:param likelihood_mask: 2d boolean array of pixels to be counted in the likelihood calculation/linear
optimization
:param psf_error_map_bool_list: list of boolean of length of point source models.
Indicates whether PSF error map is used for the point source model stated as the index.
:param kwargs_pixelbased: keyword arguments with various settings related to the pixel-based solver
(see SLITronomy documentation) being applied to the point sources.
"""
# SHAPELETS nmax must be set at initialization for linear solver
if (
source_model_class is not None
and "SHAPELETS" in source_model_class.profile_type_list
):
index = source_model_class.profile_type_list.index("SHAPELETS")
shapelet = source_model_class.func_list[index]
if not hasattr(shapelet, "num_param"):
raise ValueError(
"SHAPELETS n_max must be set at initialization. Use profile_kwargs_list when initializing LightModel.\n"
"If using kwargs_model, see documentation of source_light_profile_kwargs_list in class_creator.create_class_instances()"
)
if (
lens_light_model_class is not None
and "SHAPELETS" in lens_light_model_class.profile_type_list
):
index = lens_light_model_class.profile_type_list.index("SHAPELETS")
shapelet = lens_light_model_class.func_list[index]
if not hasattr(shapelet, "num_param"):
raise ValueError(
"SHAPELETS n_max must be set at initialization. Use profile_kwargs_list when initializing LightModel.\n"
"If using kwargs_model, see documentation of lens_light_profile_kwargs_list in class_creator.create_class_instances()"
)
super(ImageLinearFit, self).__init__(
data_class,
psf_class=psf_class,
lens_model_class=lens_model_class,
source_model_class=source_model_class,
lens_light_model_class=lens_light_model_class,
point_source_class=point_source_class,
extinction_class=extinction_class,
kwargs_numerics=kwargs_numerics,
likelihood_mask=likelihood_mask,
psf_error_map_bool_list=psf_error_map_bool_list,
kwargs_pixelbased=kwargs_pixelbased,
)
# prepare to use fft convolution for the natwt linear solver
if self.Data.likelihood_method() == "interferometry_natwt":
raise ValueError(
"interferometry_natwt linear solver not supported in jaxtronomy yet. Only diagonal likelihood method is supported."
)
# self._convolution = PixelKernelConvolution(
# kernel=self.PSF.kernel_point_source
# )
[docs]
@partial(jit, static_argnums=(0, 7))
def image_linear_solve(
self,
kwargs_lens=None,
kwargs_source=None,
kwargs_lens_light=None,
kwargs_ps=None,
kwargs_extinction=None,
kwargs_special=None,
inv_bool=False,
):
"""Computes the image (lens and source surface brightness with a given lens
model). By default, the linear parameters are computed with a weighted linear
least square optimization (i.e. flux normalization of the brightness profiles)
However in case of pixel-based modelling, pixel values are constrained by an
external solver (e.g. SLITronomy).
:param kwargs_lens: list of keyword arguments corresponding to the superposition
of different lens profiles
:param kwargs_source: list of keyword arguments corresponding to the
superposition of different source light profiles
:param kwargs_lens_light: list of keyword arguments corresponding to different
lens light surface brightness profiles
:param kwargs_ps: keyword arguments corresponding to "other" parameters, such as
external shear and point source image positions
:param kwargs_extinction: list of keyword arguments for extinction model
:param kwargs_special: list of special keyword arguments
:param inv_bool: if True, invert the full linear solver Matrix Ax = y for the
purpose of the covariance matrix. This has no impact in case of pixel-based
modelling.
:return: 2d array of surface brightness pixels of the optimal solution of the
linear parameters to match the data
"""
# NOTE: pixelbased solve is not supported in jaxtronomy
# if self._pixelbased_bool is True:
# model, model_error, cov_param, param = self.image_pixelbased_solve(
# kwargs_lens,
# kwargs_source,
# kwargs_lens_light,
# kwargs_ps,
# kwargs_extinction,
# kwargs_special,
# )
# if self.Data.likelihood_method() == "diagonal":
A = ImageLinearFit.linear_response_matrix(
self,
kwargs_lens,
kwargs_source,
kwargs_lens_light,
kwargs_ps,
kwargs_extinction,
kwargs_special,
)
C_D_response, model_error = ImageModel.error_response(
self, kwargs_lens, kwargs_ps, kwargs_special=kwargs_special
)
d = self.data_response
param, cov_param, wls_model = de_lens.get_param_WLS(
A.T, 1 / C_D_response, d, inv_bool=inv_bool
)
model = self.array_masked2image(wls_model)
_, _, _, _ = ImageLinearFit.update_linear_kwargs(
self, param, kwargs_lens, kwargs_source, kwargs_lens_light, kwargs_ps
)
# TODO: Implement this
# elif self.Data.likelihood_method() == "interferometry_natwt":
# (
# model,
# model_error,
# cov_param,
# param,
# ) = self._image_linear_solve_interferometry_natwt(
# kwargs_lens,
# kwargs_source,
# kwargs_lens_light,
# kwargs_ps,
# kwargs_extinction,
# kwargs_special,
# )
# else:
# raise ValueError(
# "likelihood_method %s not supported!" % self.Data.likelihood_method()
# )
return model, model_error, cov_param, param
# def image_pixelbased_solve(
# self,
# kwargs_lens=None,
# kwargs_source=None,
# kwargs_lens_light=None,
# kwargs_ps=None,
# kwargs_extinction=None,
# kwargs_special=None,
# init_lens_light_model=None,
# ):
# """Computes the image (lens and source surface brightness with a given lens
# model) using the pixel-based solver.
# :param kwargs_lens: list of keyword arguments corresponding to the superposition
# of different lens profiles
# :param kwargs_source: list of keyword arguments corresponding to the
# superposition of different source light profiles
# :param kwargs_lens_light: list of keyword arguments corresponding to different
# lens light surface brightness profiles
# :param kwargs_ps: keyword arguments corresponding to point sources
# :param kwargs_extinction: keyword arguments corresponding to dust extinction
# :param kwargs_special: keyword arguments corresponding to "special" parameters
# :param init_lens_light_model: optional initial guess for the lens surface
# brightness
# :return: 2d array of surface brightness pixels of the optimal solution of the
# linear parameters to match the data
# """
# _, model_error = ImageModel.error_response(
# self, kwargs_lens, kwargs_ps, kwargs_special=kwargs_special
# )
# model, param, _ = self.PixelSolver.solve(
# kwargs_lens,
# kwargs_source,
# kwargs_lens_light=kwargs_lens_light,
# kwargs_ps=kwargs_ps,
# kwargs_special=kwargs_special,
# init_lens_light_model=init_lens_light_model,
# )
# cov_param = None
# _, _ = self.update_pixel_kwargs(kwargs_source, kwargs_lens_light)
# _, _, _, _ = ImageLinearFit.update_linear_kwargs(
# self, param, kwargs_lens, kwargs_source, kwargs_lens_light, kwargs_ps
# )
# return model, model_error, cov_param, param
[docs]
@partial(jit, static_argnums=(0, 7, 9))
def likelihood_data_given_model(
self,
kwargs_lens=None,
kwargs_source=None,
kwargs_lens_light=None,
kwargs_ps=None,
kwargs_extinction=None,
kwargs_special=None,
source_marg=False,
linear_prior=None,
check_positive_flux=False,
):
"""Computes the likelihood of the data given a model This is specified with the
non-linear parameters and a linear inversion and prior marginalisation.
:param kwargs_lens: list of keyword arguments corresponding to the superposition
of different lens profiles
:param kwargs_source: list of keyword arguments corresponding to the
superposition of different source light profiles
:param kwargs_lens_light: list of keyword arguments corresponding to different
lens light surface brightness profiles
:param kwargs_ps: keyword arguments corresponding to "other" parameters, such as
external shear and point source image positions
:param kwargs_extinction: list of keyword arguments for extinction model
:param kwargs_special: list of special keyword arguments
:param source_marg: bool, performs a marginalization over the linear parameters
:param linear_prior: linear prior width in eigenvalues
:param check_positive_flux: bool, if True, checks whether the linear inversion
resulted in non-negative flux components and applies a punishment in the
likelihood if so.
:return: log likelihood (natural logarithm), linear parameter list
"""
im_sim, model_error, cov_matrix, param = ImageLinearFit.image_linear_solve(
self,
kwargs_lens,
kwargs_source,
kwargs_lens_light,
kwargs_ps,
kwargs_extinction,
kwargs_special,
inv_bool=source_marg,
)
# compute X^2
logL = self.likelihood_data_given_model_solution(
im_sim,
model_error,
cov_matrix,
param,
kwargs_lens,
kwargs_source,
kwargs_lens_light,
kwargs_ps,
source_marg=source_marg,
linear_prior=linear_prior,
check_positive_flux=check_positive_flux,
)
return logL, param
[docs]
@partial(jit, static_argnums=(0, 9, 11))
def likelihood_data_given_model_solution(
self,
model,
model_error,
cov_matrix,
param,
kwargs_lens,
kwargs_source,
kwargs_lens_light,
kwargs_ps,
source_marg=False,
linear_prior=None,
check_positive_flux=False,
):
"""
:param model: 2d array, image model
:param model_error: 2d array, uncertainties in each pixel
:param cov_matrix: 2d array, covariance matrix
:param param: linear parameter vector corresponding to the response matrix
:param kwargs_lens: list of dicts containing lens model keyword arguments
:param kwargs_source: list of dicts containing source model keyword arguments
:param kwargs_lens_light: list of dicts containing lens light model keyword arguments
:param kwargs_ps: list of dicts containing point source keyword arguments
:param kwargs_extinction: list of keyword arguments for extinction model
:param kwargs_special: list of special keyword arguments
:param source_marg: bool, performs a marginalization over the linear parameters
:param linear_prior: linear prior width in eigenvalues
:param check_positive_flux: bool, if True, checks whether the linear inversion
resulted in non-negative flux components and applies a punishment in the
likelihood if so.
:return: float, likelihood data given model
"""
logL = self.Data.log_likelihood(model, self.likelihood_mask, model_error)
if self._pixelbased_bool is False:
if cov_matrix is not None and source_marg:
marg_const = de_lens.marginalization_new(
cov_matrix, d_prior=linear_prior
)
logL += marg_const
if check_positive_flux is True:
_, _, _, _ = ImageLinearFit.update_linear_kwargs(
self, param, kwargs_lens, kwargs_source, kwargs_lens_light, kwargs_ps
)
bool_ = self.check_positive_flux(
kwargs_source, kwargs_lens_light, kwargs_ps
)
logL = jnp.where(bool_, logL, logL - 10**8)
return logL
[docs]
@partial(jit, static_argnums=(0,))
def num_param_linear(
self, kwargs_lens, kwargs_source, kwargs_lens_light, kwargs_ps
):
"""
:param kwargs_lens: list of dicts containing lens model keyword arguments
:param kwargs_source: list of dicts containing source model keyword arguments
:param kwargs_lens_light: list of dicts containing lens light model keyword arguments
:param kwargs_ps: list of dicts containing point source keyword arguments
:return: number of linear coefficients to be solved for in the linear inversion
"""
num = 0
if self._pixelbased_bool is False:
num += self.SourceModel.num_param_linear(kwargs_source)
num += self.LensLightModel.num_param_linear(kwargs_lens_light)
num += self.PointSource.num_basis(kwargs_ps, kwargs_lens)
return num
[docs]
@partial(jit, static_argnums=(0, 7))
def linear_response_matrix(
self,
kwargs_lens,
kwargs_source,
kwargs_lens_light,
kwargs_ps,
kwargs_extinction=None,
kwargs_special=None,
unconvolved=False,
):
"""Computes the linear response matrix (m x n), with n being the data size and m
being the coefficients.
The calculation is done by
- first (optional) computing differential extinctions
- adding linear components of the lensed source(s)
- adding linear components of the unlensed components (i.e. deflector)
- adding point sources (can be multiple lensed or stars in the field)
:param kwargs_lens: list of keyword arguments corresponding to the superposition of different lens profiles
:param kwargs_source: list of keyword arguments corresponding to the superposition of different source light profiles
:param kwargs_lens_light: list of keyword arguments corresponding to different lens light surface brightness profiles
:param kwargs_ps: keyword arguments corresponding to "other" parameters, such as external shear and point source image positions
:param kwargs_extinction: list of keyword arguments for extinction model
:param kwargs_special: list of special keyword arguments
:param unconvolved: bool, if True, computes components without convolution kernel (will not work for point sources)
:return: response matrix (m x n)
"""
x_grid, y_grid = self.ImageNumerics.coordinates_evaluate
source_light_response, _ = self.source_mapping.image_flux_split(
x_grid, y_grid, kwargs_lens, kwargs_source, kwargs_special
)
n_source = len(source_light_response)
# extinction = self._extinction.extinction(
# x_grid,
# y_grid,
# kwargs_extinction=kwargs_extinction,
# kwargs_special=kwargs_special,
# )
lens_light_response, _ = self.LensLightModel.functions_split(
x_grid, y_grid, kwargs_lens_light
)
n_lens_light = len(lens_light_response)
ra_pos, dec_pos, amp, _ = self.point_source_linear_response_set(
kwargs_ps, kwargs_lens, kwargs_special, with_amp=False
)
n_points = len(ra_pos)
num_param = n_points + n_lens_light + n_source
num_response = self.num_data_evaluate
A = jnp.zeros((num_param, num_response))
# This function loops through the light responses
def body_fun(i, val):
A, n, light_response = val
image = light_response.at[i].get()
# NOTE: Primary beam not supported in jaxtronomy
# multiply with primary beam before convolution
# if self._pb is not None:
# image *= self._pb_1d
# image *= extinction
image = self.ImageNumerics.re_size_convolve(image, unconvolved=unconvolved)
A = A.at[i + n].set(jnp.nan_to_num(self.image2array_masked(image)))
return A, n, light_response
# response of lensed source profile
# this if-statement is needed to prevent compiler from trying to index into an empty array
if n_source != 0:
A, _, _ = lax.fori_loop(
0, n_source, body_fun, (A, 0, source_light_response)
)
# response of deflector light profile (or any other un-lensed extended components)
if n_lens_light != 0:
A, _, _ = lax.fori_loop(
0, n_lens_light, body_fun, (A, n_source, lens_light_response)
)
# response of point sources
n = n_source + n_lens_light
for i in range(0, n_points):
# NOTE: Primary beam not supported in jaxtronomy
# raise warnings when primary beam is attempted to be applied for point sources
# if self._pb is not None:
# raise Warning("Antenna primary beam does not apply to point sources!")
image = self.ImageNumerics.point_source_rendering(
ra_pos[i], dec_pos[i], amp[i]
)
A = A.at[i + n].set(jnp.nan_to_num(self.image2array_masked(image)))
return A * self._flux_scaling
[docs]
@partial(jit, static_argnums=(0,))
def update_linear_kwargs(
self, param, kwargs_lens, kwargs_source, kwargs_lens_light, kwargs_ps
):
"""Links linear parameters to kwargs arguments.
:param param: linear parameter vector corresponding to the response matrix
:param kwargs_lens: list of dicts containing lens model keyword arguments
:param kwargs_source: list of dicts containing source model keyword arguments
:param kwargs_lens_light: list of dicts containing lens light model keyword
arguments
:param kwargs_ps: list of dicts containing point source keyword arguments
:return: updated list of kwargs with linear parameter values
"""
i = 0
param = jnp.asarray(param)
kwargs_source, i = self.SourceModel.update_linear(
param, i, kwargs_list=kwargs_source
)
kwargs_lens_light, i = self.LensLightModel.update_linear(
param, i, kwargs_list=kwargs_lens_light
)
kwargs_ps, i = self.PointSource.update_linear(param, i, kwargs_ps, kwargs_lens)
return kwargs_lens, kwargs_source, kwargs_lens_light, kwargs_ps
[docs]
@partial(jit, static_argnums=0)
def linear_param_from_kwargs(self, kwargs_source, kwargs_lens_light, kwargs_ps):
"""Returns list of the linear amplitudes from the keyword arguments.
:param kwargs_source: list of dicts containing source model keyword arguments
:param kwargs_lens_light: list of dicts containing lens light model keyword
arguments
:param kwargs_ps: list of dicts containing point source keyword arguments
:return: list of linear coefficients
"""
param = []
param += self.SourceModel.linear_param_from_kwargs(kwargs_source)
param += self.LensLightModel.linear_param_from_kwargs(kwargs_lens_light)
param += self.PointSource.linear_param_from_kwargs(kwargs_ps)
return param
# NOTE: pixelbased profiles not supported in jaxtronomy
# def update_pixel_kwargs(self, kwargs_source, kwargs_lens_light):
# """Update kwargs arguments for pixel-based profiles with fixed properties such
# as their number of pixels, scale, and center coordinates (fixed to the origin).
# :param kwargs_source: list of keyword arguments corresponding to the
# superposition of different source light profiles
# :param kwargs_lens_light: list of keyword arguments corresponding to the
# superposition of different lens light profiles
# :return: updated kwargs_source and kwargs_lens_light
# """
# # in case the source plane grid size has changed, update the kwargs accordingly
# ss_factor_source = self.SourceNumerics.grid_supersampling_factor
# kwargs_source[0]["n_pixels"] = int(
# self.Data.num_pixel * ss_factor_source**2
# ) # effective number of pixels in source plane
# kwargs_source[0]["scale"] = (
# self.Data.pixel_width / ss_factor_source
# ) # effective pixel size of source plane grid
# # pixelated reconstructions have no well-defined center, we put it arbitrarily at (0, 0), center of the image
# kwargs_source[0]["center_x"] = 0
# kwargs_source[0]["center_y"] = 0
# # do the same if the lens light has been reconstructed
# if kwargs_lens_light is not None and len(kwargs_lens_light) > 0:
# kwargs_lens_light[0]["n_pixels"] = self.Data.num_pixel
# kwargs_lens_light[0]["scale"] = self.Data.pixel_width
# kwargs_lens_light[0]["center_x"] = 0
# kwargs_lens_light[0]["center_y"] = 0
# return kwargs_source, kwargs_lens_light
[docs]
@partial(jit, static_argnums=(0,))
def error_map_source(self, kwargs_source, x_grid, y_grid, cov_param):
"""Variance of the linear source reconstruction in the source plane coordinates,
computed by the diagonal elements of the covariance matrix of the source
reconstruction as a sum of the errors of the basis set.
:param kwargs_source: keyword arguments of source model
:param x_grid: x-axis of positions to compute error map
:param y_grid: y-axis of positions to compute error map
:param cov_param: covariance matrix of linear inversion parameters
:return: diagonal covariance errors at the positions (x_grid, y_grid)
"""
error_map = jnp.zeros_like(x_grid)
basis_functions, _ = self.SourceModel.functions_split(
x_grid, y_grid, kwargs_source
)
n_source = len(basis_functions)
basis_functions = jnp.array(basis_functions)
if cov_param is not None:
for i in range(len(error_map)):
error_map = error_map.at[i].set(
basis_functions[:, i]
.T.dot(cov_param[:n_source, :n_source])
.dot(basis_functions[:, i])
)
return error_map
[docs]
@partial(jit, static_argnums=(0, 4))
def point_source_linear_response_set(
self, kwargs_ps, kwargs_lens, kwargs_special=None, with_amp=True
):
"""
:param kwargs_ps: point source keyword argument list
:param kwargs_lens: lens model keyword argument list
:param kwargs_special: special keyword argument list, may include 'delta_x_image' and 'delta_y_image'
:param with_amp: bool, if True, relative magnification between multiply imaged point sources are held fixed.
:return: list of positions and amplitudes split in different basis components with applied astrometric corrections
"""
ra_pos, dec_pos, amp, _ = self.PointSource.linear_response_set(
kwargs_ps, kwargs_lens, with_amp=with_amp
)
n_points = len(ra_pos)
if kwargs_special is not None:
if "delta_x_image" in kwargs_special:
delta_x, delta_y = (
kwargs_special["delta_x_image"],
kwargs_special["delta_y_image"],
)
k = 0
n = len(delta_x)
for i in range(n_points):
for j in range(len(ra_pos[i])):
if k >= n:
break
ra_pos[i] = (
ra_pos[i].at[j].set(ra_pos[i].at[j].get() + delta_x[k])
)
dec_pos[i] = (
dec_pos[i].at[j].set(dec_pos[i].at[j].get() + delta_y[k])
)
k += 1
return ra_pos, dec_pos, amp, n_points
[docs]
@partial(jit, static_argnums=0)
def check_positive_flux(self, kwargs_source, kwargs_lens_light, kwargs_ps):
"""Checks whether the surface brightness profiles contain positive fluxes and
returns bool if True.
:param kwargs_source: source surface brightness keyword argument list
:param kwargs_lens_light: lens surface brightness keyword argument list
:param kwargs_ps: point source keyword argument list
:return: boolean
"""
pos_bool = True
pos_bool_ps = self.PointSource.check_positive_flux(kwargs_ps)
# if self._pixelbased_bool is True:
# # this constraint must be handled by the pixel-based solver
# pos_bool_source = True
# pos_bool_lens_light = True
# else:
pos_bool_source = self.SourceModel.check_positive_flux_profile(kwargs_source)
pos_bool_lens_light = self.LensLightModel.check_positive_flux_profile(
kwargs_lens_light
)
pos_bool = jnp.where(pos_bool_ps, pos_bool, False)
pos_bool = jnp.where(pos_bool_source, pos_bool, False)
pos_bool = jnp.where(pos_bool_lens_light, pos_bool, False)
return pos_bool
# linear solver for interferometric natwt method
# def _image_linear_solve_interferometry_natwt(
# self,
# kwargs_lens=None,
# kwargs_source=None,
# kwargs_lens_light=None,
# kwargs_ps=None,
# kwargs_extinction=None,
# kwargs_special=None,
# ):
# """'interferometry_natwt' method does NOT support model_error, cov_param. The
# interferometry linear solver just does the linear solving to get the optimal
# linear amplitudes and apply the marginalized amplitudes to make the model
# images.
# :param kwargs_lens: list of dicts containing lens model keyword arguments
# :param kwargs_source: list of dicts containing source model keyword arguments
# :param kwargs_lens_light: list of dicts containing lens light model keyword arguments
# :param kwargs_ps: list of dicts containing point source keyword arguments
# :param kwargs_extinction: list of keyword arguments for extinction model
# :param kwargs_special: list of special keyword arguments
# :return: model, model_error, cov_param, param
# model and param are the same returns of self._image_linear_solve_interferometry_natwt_solving(A, d) function
# model_error =0 and cov_param = None for the interferometric method.
# """
# A = ImageLinearFit.linear_response_matrix(
# self,
# kwargs_lens,
# kwargs_source,
# kwargs_lens_light,
# kwargs_ps,
# kwargs_extinction,
# kwargs_special,
# unconvolved=True,
# )
# d = self.data_response
# model, param = self._image_linear_solve_interferometry_natwt_solving(A, d)
# model_error = 0 # just a place holder
# cov_param = None # just a place holder
# _, _, _, _ = ImageLinearFit.update_linear_kwargs(
# self, param, kwargs_lens, kwargs_source, kwargs_lens_light, kwargs_ps
# )
# return model, model_error, cov_param, param
# def _image_linear_solve_interferometry_natwt_solving(self, A, d):
# """Linearly solve the amplitude of each light profile response to the natural
# weighting interferometry images, based on (placeholder for Nan Zhang's paper).
# Theories:
# Suppose there are a set of light responses :math:`\\{x_i\\}`, we want to solve the set of amplitudes :math:`\\{\\alpha_i\\}`,
# such that minimizes the chi^2 given by
# .. math::
# \\chi^2 = (d - A_{PSF}\\sum_i \\alpha_i x_i)^TC^{-1}(d - A_{PSF}\\sum_i \\alpha_i x_i),
# where :math:`A_{PSF}` is the PSF convolution operation matrix (not to be confused with the input A of this function)
# and :math:`C` is the noise covariance matrix. :math:`d` is the data image.
# For natural weighting interferometric images, we have :math:`C = \\sigma^2 A_{PSF}`,
# (see Section 3.2 of https://doi.org/10.1093/mnras/staa2740 for the relation of natural weighting covariance matrix and PSF convolution)
# therefore the chi^2 function simplifies to
# .. math::
# \\chi^2 = \\frac{1}{\\sigma^2}(d^TA_{PSF}^{-1}d + \\sum_{i,j}\\alpha_i\\alpha_j x_i^TA_{PSF}x_j - 2\\sum_{i}x_i^Td),
# from which the optimal amplitudes :math:`\\{\\alpha_i\\}` can be solved linearly by solving
# .. math::
# \\sum_{j} M_{ij}\\alpha_{j} = b_i,
# where :math:`M_{ij} = \\frac{1}{\\sigma^2}x_i^TA_{PSF}x_j` and :math:`b_{i} = \\frac{1}{\\sigma^2}x_i^Td`.
# The steps of this function are:
# (1.) Making the entries :math:`M_{ij}` and :math:`b_i` defined above.
# (2.) Solve the linear function to get the optimal amplitudes.
# (3.) Apply these optimal amplitudes to make unconvolved and convolved model images.
# The output model images are in the form [array1, array2].
# (Note that this is different from the non-interferometric linear solver of Lenstronomy,
# this output form saves time for likelihood computations in imaging_data for interferometric method.)
# array1 is the unconvolved model image :math:`array1 = \\sum_i \\alpha_i x_i`, where :math:`\\alpha_i` is the solved optimal amplitudes.
# array2 is the convolved model image :math:`array2 = A_{PSF}\\sum_i \\alpha_i x_i`, where :math:`\\alpha_i`.
# :param A: response of unconvolved light profiles, [x_1, x_2, ...]
# :param d: data image, d
# :return: [array1, array2], [amp_array]
# where the [array1, array2] are unconvolved and convolved model images with solved amplitudes
# and [amp_array] are the solved optimal amplitudes.
# """
# num_of_light, num_of_image_pixel = np.shape(A)
# A_convolved = np.zeros(np.shape(A))
# # convolve each response separately
# for i in range(num_of_light):
# A_convolved[i] = util.image2array(
# self._convolution._static_fft(util.array2image(A[i]), mode="same")
# )
# M = np.zeros((num_of_light, num_of_light))
# for i in range(num_of_light):
# for j in range(num_of_light):
# if j < i:
# M[i, j] = M[j, i]
# else:
# M[i, j] = np.sum(A_convolved[j] * A[i])
# b = np.zeros((num_of_light))
# for i in range(num_of_light):
# b[i] = np.sum(A[i] * (d))
# param_amps = np.linalg.lstsq(M, b, rcond=None)[0]
# clean_temp = np.zeros((num_of_image_pixel))
# dirty_temp = np.zeros((num_of_image_pixel))
# for i in range(num_of_light):
# clean_temp += param_amps[i] * A[i]
# dirty_temp += param_amps[i] * A_convolved[i]
# clean_model = util.array2image(clean_temp)
# dirty_model = util.array2image(dirty_temp)
# model = [clean_model, dirty_model]
# return model, param_amps