from lenstronomy.ImSim.image_model import ImageModel
from lenstronomy.LightModel.light_model import LightModel
from lenstronomy.LensModel.lens_model import LensModel
from lenstronomy.ImSim.image2source_mapping import Image2SourceMapping
from lenstronomy.Util import util
import numpy as np
[docs]
class TracerModelSource(ImageModel):
"""Tracer model class, inherits ImageModel."""
[docs]
def __init__(
self,
data_class,
psf_class=None,
lens_model_class=None,
source_model_class=None,
lens_light_model_class=None,
point_source_class=None,
extinction_class=None,
tracer_source_class=None,
kwargs_numerics=None,
likelihood_mask=None,
psf_error_map_bool_list=None,
kwargs_pixelbased=None,
tracer_partition=None,
tracer_type="LINEAR",
):
"""
: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 tracer_source_class: LightModel() instance describing the tracers of the source
: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.
:param tracer_partition: in case of tracer models for specific sub-parts of the surface brightness model
[[list of light profiles, list of tracer profiles], [list of light profiles, list of tracer profiles], [...], ...]
:type tracer_partition: None or list
:param tracer_type: LINEAR or LOG. If the tracer is in log units, it is converted to linear units, summed,
and then converted back to log units.
:type tracer_partition: str
"""
if likelihood_mask is None:
likelihood_mask = np.ones_like(data_class.data)
self.likelihood_mask = np.array(likelihood_mask, dtype=bool)
self._mask1d = util.image2array(self.likelihood_mask)
if tracer_partition is None:
tracer_partition = [[None, None]]
self._tracer_partition = tracer_partition
if tracer_type not in ["LINEAR", "LOG"]:
raise Exception(
"Unknown input tracer_type: {0}. Only two tracer types are currently supported: LINEAR and LOG. Please convert your tracer to linear/log units.".format(
tracer_type
)
)
self._tracer_type = tracer_type
super(TracerModelSource, 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,
kwargs_pixelbased=kwargs_pixelbased,
)
if psf_error_map_bool_list is None:
psf_error_map_bool_list = [True] * len(
self.PointSource.point_source_type_list
)
self._psf_error_map_bool_list = psf_error_map_bool_list
if tracer_source_class is None:
tracer_source_class = LightModel(light_model_list=[])
if lens_model_class is None:
lens_model_class = LensModel(lens_model_list=[])
self.tracer_mapping = Image2SourceMapping(
lens_model=lens_model_class, source_model=tracer_source_class
)
self.tracer_source_class = tracer_source_class
[docs]
def tracer_model(
self,
kwargs_tracer_source,
kwargs_lens,
kwargs_source,
kwargs_extinction=None,
kwargs_special=None,
de_lensed=False,
):
"""Tracer model as a convolved surface brightness weighted quantity conv(tracer
* surface brightness) / conv(surface brightness)
:param kwargs_tracer_source:
:param kwargs_lens:
:param kwargs_source:
:return: model predicted observed tracer component
"""
tracer_brightness_conv = np.zeros_like(self.Data.data)
source_light_conv = np.zeros_like(self.Data.data)
for [k_light, k_tracer] in self._tracer_partition:
source_light_k = self._source_surface_brightness_analytical_numerics(
kwargs_source,
kwargs_lens,
kwargs_extinction,
kwargs_special=kwargs_special,
de_lensed=de_lensed,
k=k_light,
)
source_light_conv_k = self.ImageNumerics.re_size_convolve(
source_light_k, unconvolved=False
)
source_light_conv_k[source_light_conv_k < 10 ** (-20)] = 10 ** (-20)
tracer_k = self._tracer_model_source(
kwargs_tracer_source, kwargs_lens, de_lensed=de_lensed, k=k_tracer
)
if self._tracer_type == "LINEAR":
tracer_brightness_conv_k = self.ImageNumerics.re_size_convolve(
tracer_k * source_light_k, unconvolved=False
)
tracer_brightness_conv += tracer_brightness_conv_k
if self._tracer_type == "LOG":
lin_tracer_k = 10 ** (tracer_k)
lin_tracer_brightness_conv_k = self.ImageNumerics.re_size_convolve(
lin_tracer_k * source_light_k, unconvolved=False
)
tracer_brightness_conv += lin_tracer_brightness_conv_k
source_light_conv += source_light_conv_k
if self._tracer_type == "LINEAR":
return tracer_brightness_conv / source_light_conv
if self._tracer_type == "LOG":
return np.log10(tracer_brightness_conv / source_light_conv)
def _tracer_model_source(
self, kwargs_tracer_source, kwargs_lens, de_lensed=False, k=None
):
"""
:param kwargs_tracer_source:
:param kwargs_lens:
:return:
"""
ra_grid, dec_grid = self.ImageNumerics.coordinates_evaluate
if de_lensed is True:
source_light = self.tracer_source_class.surface_brightness(
ra_grid, dec_grid, kwargs_tracer_source, k=k
)
else:
source_light = self.tracer_mapping.image_flux_joint(
ra_grid, dec_grid, kwargs_lens, kwargs_tracer_source, k=k
)
return source_light
[docs]
def likelihood_data_given_model(
self,
kwargs_tracer_source,
kwargs_lens,
kwargs_source,
kwargs_extinction=None,
kwargs_special=None,
):
model = self.tracer_model(
kwargs_tracer_source,
kwargs_lens,
kwargs_source,
kwargs_extinction,
kwargs_special,
)
log_likelihood = self.Data.log_likelihood(
model, self.likelihood_mask, additional_error_map=0
)
return log_likelihood
@property
def num_data_evaluate(self):
"""Number of data points to be used in the linear solver :return:"""
return int(np.sum(self.likelihood_mask))