Repository: lenstronomy/lenstronomy

Issue:
Time-delay likelihood is incorrect if several point sources are used. 
This statement is always False when self._measurement_bool_list is a list of lists, which is the case when you have several 'LENSED_POSITION' point sources: https://github.com/lenstronomy/lenstronomy/blob/f731c7c9fe2fd91ba6921184baaf151074a02ea5/lenstronomy/Sampling/Likelihoods/time_delay_likelihood.py#L68

I am working on a PR to solve this.

Discussion:
@sibirrer: @martin-millon the original use of this is intent to only have one True/False statement per point_source_model, either a source is time-variable and measurements exist, or the source is static. We either need to specify the use or indeed fix it to not cause confusions. Thank you for pointing this out and looking forward to the PR!@martin-millon: The proposed solution should conserve the previous behavior if default arguments are provided. You can now provide a list of list of boolean in `time_delay_measurement_bool_list`. For example, for 2-point sources, imaged 4 times each (3 time-delays), you can provide `time_delay_measurement_bool_list=[[True, True, True], [True, False, False]]`. @martin-millon: Fixed in #680 

Solution patch:
diff --git a/lenstronomy/ImSim/image2source_mapping.py b/lenstronomy/ImSim/image2source_mapping.py
index df606dee..3c8e0f8c 100644
--- a/lenstronomy/ImSim/image2source_mapping.py
+++ b/lenstronomy/ImSim/image2source_mapping.py
@@ -75,11 +75,12 @@ class Image2SourceMapping(object):
                     "length of redshift_list must correspond to length of light_model_list"
                 )
 
-            if np.max(self._source_redshift_list) > self._lens_model.z_source:
-                raise ValueError(
-                    "redshift of source_redshift_list have to be smaller or equal to "
-                    "the one specified in the lens model."
-                )
+            if len(self._source_redshift_list) > 0:
+                if np.max(self._source_redshift_list) > self._lens_model.z_source:
+                    raise ValueError(
+                        "redshift of source_redshift_list have to be smaller or equal to "
+                        "the one specified in the lens model."
+                    )
 
             # turn off multi source plane if all sources are at the same redshift
             if len(list(set(self._source_redshift_list))) == 1:
diff --git a/lenstronomy/LensModel/lens_model.py b/lenstronomy/LensModel/lens_model.py
index 7e2b7039..bb8602a4 100644
--- a/lenstronomy/LensModel/lens_model.py
+++ b/lenstronomy/LensModel/lens_model.py
@@ -26,6 +26,7 @@ class LensModel(object):
         numerical_alpha_class=None,
         observed_convention_index=None,
         z_source_convention=None,
+        z_lens_convention=None,
         cosmo_interp=False,
         z_interp_stop=None,
         num_z_interp=100,
@@ -55,6 +56,9 @@ class LensModel(object):
          positions. The code will compute the physical locations when performing computations
         :param z_source_convention: float, redshift of a source to define the reduced deflection angles of the lens
          models. If None, 'z_source' is used.
+        :param z_lens_convention: float, redshift of a lens plane to define the
+         effective time-delay distance. Only needed if distance ratios are
+         sampled. If None, the first lens redshift is used.
         :param cosmo_interp: boolean (only employed in multi-plane mode), interpolates astropy.cosmology distances for
          faster calls when accessing several lensing planes
         :param z_interp_stop: (only in multi-plane with cosmo_interp=True); maximum redshift for distance interpolation
@@ -142,6 +146,7 @@ class LensModel(object):
                     numerical_alpha_class=numerical_alpha_class,
                     observed_convention_index=observed_convention_index,
                     z_source_convention=z_source_convention,
+                    z_lens_convention=z_lens_convention,
                     cosmo_interp=cosmo_interp,
                     z_interp_stop=z_interp_stop,
                     num_z_interp=num_z_interp,
diff --git a/lenstronomy/LightModel/Profiles/shapelets.py b/lenstronomy/LightModel/Profiles/shapelets.py
index 6887efb1..6dcd2295 100644
--- a/lenstronomy/LightModel/Profiles/shapelets.py
+++ b/lenstronomy/LightModel/Profiles/shapelets.py
@@ -34,7 +34,7 @@ class Shapelets(object):
     The two-dimensional basis function is
 
     .. math::
-        \\phi_{\\bf n}({\bf x}) \\equiv \\phi_{n1}(x1) \\phi_{n2}(x2)
+        \\phi_{\\bf n}({\\bf x}) \\equiv \\phi_{n1}(x1) \\phi_{n2}(x2)
 
     where :math:`{\\bf n} \\equiv (n1, n2)` and :math:`{\\bf x} \\equiv (x1, x2)`.
 
diff --git a/lenstronomy/Plots/lens_plot.py b/lenstronomy/Plots/lens_plot.py
index 8bbb1a5b..1bd2c751 100644
--- a/lenstronomy/Plots/lens_plot.py
+++ b/lenstronomy/Plots/lens_plot.py
@@ -149,6 +149,9 @@ def convergence_plot(
     vmin=-1,
     vmax=1,
     cmap="Greys",
+    with_color_bar=False,
+    colorbar_label=r"$\log_{10}(\kappa)$",
+    font_size=20,
     **kwargs,
 ):
     """Plot convergence.
@@ -163,6 +166,9 @@ def convergence_plot(
     :param vmax: matplotlib vmax
     :param cmap: matplotlib cmap
     :param kwargs: keyword arguments for matshow
+    :param with_color_bar: bool, if True, shows color bar
+    :param colorbar_label: string, label of color bar
+    :param font_size: int, font size of color bar label
     :return: matplotlib axis instance with convergence plot
     """
     x_grid, y_grid = pixel_grid.pixel_coordinates
@@ -170,7 +176,7 @@ def convergence_plot(
     y_grid1d = util.image2array(y_grid)
     kappa_result = lens_model.kappa(x_grid1d, y_grid1d, kwargs_lens)
     kappa_result = util.array2image(kappa_result)
-    _ = ax.matshow(
+    im = ax.matshow(
         np.log10(kappa_result),
         origin="lower",
         extent=extent,
@@ -179,6 +185,12 @@ def convergence_plot(
         vmax=vmax,
         **kwargs,
     )
+    if with_color_bar:
+        divider = make_axes_locatable(ax)
+        cax = divider.append_axes("right", size="5%", pad=0.05)
+        cb = plt.colorbar(im, cax=cax)
+        cb.set_label(colorbar_label, fontsize=font_size)
+
     return ax
 
 
@@ -284,6 +296,8 @@ def point_source_plot(
     source_y,
     name_list=None,
     index=None,
+    solver_type="lenstronomy",
+    kwargs_solver={},
     color="k",
     **kwargs,
 ):
@@ -302,6 +316,8 @@ def point_source_plot(
     :param name_list: list of strings, longer or equal the number of point sources. If changing this parameter, input as name_list=[[...], [...]]
     :param index: number of sources, an integer number. Default None.
     :param color: string representing the color for the source's images. Default "k".
+    :param solver_type: string, type of solver to find the image positions ('lenstronomy', 'analytical' or 'stochastic')
+    :param kwargs_solver: keyword arguments for the solver
     :param kwargs: additional plotting keyword arguments
     :return: matplotlib axis instance with figure
     """
@@ -342,6 +358,8 @@ def point_source_plot(
         x_center=x_center,
         y_center=y_center,
         min_distance=pixel_grid.pixel_width,
+        solver=solver_type,
+        **kwargs_solver,
     )
     mag_images = lens_model.magnification(theta_x, theta_y, kwargs_lens)
 
diff --git a/lenstronomy/Sampling/Likelihoods/time_delay_likelihood.py b/lenstronomy/Sampling/Likelihoods/time_delay_likelihood.py
index cb81a91f..17f656dc 100644
--- a/lenstronomy/Sampling/Likelihoods/time_delay_likelihood.py
+++ b/lenstronomy/Sampling/Likelihoods/time_delay_likelihood.py
@@ -23,7 +23,9 @@ class TimeDelayLikelihood(object):
         :param lens_model_class: instance of the LensModel() class
         :param point_source_class: instance of the PointSource() class, note: the first point source type is the one the
          time delays are imposed on
-        :param time_delay_measurement_bool_list: list of bool to indicate for which point source model a measurement is available
+        :param time_delay_measurement_bool_list: list of list of bool to indicate for which point source model a measurement is available.
+         This list must have the same length as time_delays_measured and time_delays_uncertainties.
+        Example for two point sources, imaged 4 times each: [[True, False, True], [True, True, True]]
         """
 
         if time_delays_measured is None:
@@ -49,7 +51,43 @@ class TimeDelayLikelihood(object):
                 self._delays_errors.append(np.array(time_delays_uncertainties[i]))
 
         if time_delay_measurement_bool_list is None:
-            time_delay_measurement_bool_list = [True] * self._num_point_sources
+            if self._num_point_sources == 1:
+                time_delay_measurement_bool_list = [[True] * len(time_delays_measured)]
+            else:
+                time_delay_measurement_bool_list = []
+                for i in range(self._num_point_sources):
+                    time_delay_measurement_bool_list.append(
+                        [True] * len(time_delays_measured[i])
+                    )
+        else:
+            if len(time_delay_measurement_bool_list) != self._num_point_sources:
+                raise ValueError(
+                    "time_delay_measurement_bool_list must have the same length as the number of point sources."
+                )
+            for i in range(self._num_point_sources):
+                if isinstance(
+                    time_delay_measurement_bool_list[i], (bool, np.bool_, int)
+                ):
+                    print(
+                        "Warning: time_delay_measurement_bool_list is a single bool, converting to list of bools, assuming all time delays are measured."
+                    )
+                    time_delay_measurement_bool_list[i] = [
+                        bool(time_delay_measurement_bool_list[i])
+                    ] * len(self._delays_measured[i])
+                elif isinstance(
+                    time_delay_measurement_bool_list[i], (list, np.ndarray)
+                ):
+                    if len(time_delay_measurement_bool_list[i]) != len(
+                        self._delays_measured[i]
+                    ):
+                        raise ValueError(
+                            "time_delay_measurement_bool_list and time_delays_measured need to have the same length."
+                        )
+                else:
+                    raise ValueError(
+                        "time_delay_measurement_bool_list must be a list of bools or a list of lists of bools."
+                    )
+
         self._measurement_bool_list = time_delay_measurement_bool_list
 
     def logL(self, kwargs_lens, kwargs_ps, kwargs_cosmo):
@@ -65,7 +103,8 @@ class TimeDelayLikelihood(object):
         )
         logL = 0
         for i in range(self._num_point_sources):
-            if self._measurement_bool_list[i] is True:
+            mask = np.array(self._measurement_bool_list[i])
+            if np.any(mask):
                 x_pos_, y_pos_ = x_pos[i], y_pos[i]
                 self._lensModel.change_source_redshift(
                     z_source=self._pointSource._redshift_list[i]
@@ -76,9 +115,25 @@ class TimeDelayLikelihood(object):
                 D_dt_model = kwargs_cosmo["D_dt"]
                 Ddt_scaled = self._lensModel.ddt_scaling * D_dt_model
                 delay_days = const.delay_arcsec2days(delay_arcsec, Ddt_scaled)
-                logL += self._logL_delays(
-                    delay_days, self._delays_measured[i], self._delays_errors[i]
-                )
+                mask_full = np.concatenate(
+                    ([True], mask)
+                )  # add the first image to the mask
+                if len(delay_days) - 1 != len(self._delays_measured[i]):
+                    logL += -(10**15)
+                else:
+                    if self._delays_errors[i].ndim == 1:
+                        logL += self._logL_delays(
+                            delay_days[mask_full],
+                            self._delays_measured[i][mask],
+                            self._delays_errors[i][mask],
+                        )
+                    elif self._delays_errors[i].ndim == 2:
+                        # mask the covariance matrix
+                        logL += self._logL_delays(
+                            delay_days[mask_full],
+                            self._delays_measured[i][mask],
+                            self._delays_errors[i][mask, :][:, mask],
+                        )
         return logL
 
     @staticmethod
diff --git a/lenstronomy/Sampling/likelihood.py b/lenstronomy/Sampling/likelihood.py
index 6debec2b..2386df75 100644
--- a/lenstronomy/Sampling/likelihood.py
+++ b/lenstronomy/Sampling/likelihood.py
@@ -46,6 +46,7 @@ class LikelihoodModule(object):
         max_num_images=None,
         bands_compute=None,
         time_delay_likelihood=False,
+        time_delay_measurement_bool_list=None,
         image_likelihood_mask_list=None,
         flux_ratio_likelihood=False,
         kwargs_flux_compute=None,
@@ -110,6 +111,8 @@ class LikelihoodModule(object):
             which "band" to include in the fitting
         :param time_delay_likelihood: bool, if True computes the time-delay likelihood
             of the FIRST point source
+        :param time_delay_measurement_bool_list: list of bool to indicate for which
+            point source model a measurement is available
         :param kwargs_flux_compute: keyword arguments of how to compute the image
             position fluxes (see FluxRatioLikeliood)
         :param custom_logL_addition: a definition taking as arguments (kwargs_lens,
@@ -173,6 +176,7 @@ class LikelihoodModule(object):
         self._kwargs_time_delay = {
             "time_delays_measured": time_delays_measured,
             "time_delays_uncertainties": time_delays_uncertainties,
+            "time_delay_measurement_bool_list": time_delay_measurement_bool_list,
         }
         self._kwargs_image_likelihood = {
             "source_marg": source_marg,


Test patch:
diff --git a/test/test_Sampling/test_Likelihoods/test_time_delay_likelihood.py b/test/test_Sampling/test_Likelihoods/test_time_delay_likelihood.py
index 2ab48359..4ef93d8d 100644
--- a/test/test_Sampling/test_Likelihoods/test_time_delay_likelihood.py
+++ b/test/test_Sampling/test_Likelihoods/test_time_delay_likelihood.py
@@ -10,38 +10,37 @@ import pytest
 import copy
 
 
-class TestImageLikelihood(object):
+class TestTimeDelayLikelihood(object):
     def setup_method(self):
-        pass
-
-    def test_logL(self):
-        z_source = 1.5
-        z_lens = 0.5
-        cosmo = FlatLambdaCDM(H0=70, Om0=0.3, Ob0=0.0)
-        lensCosmo = LensCosmo(cosmo=cosmo, z_lens=z_lens, z_source=z_source)
+        self.z_source = 1.5
+        self.z_lens = 0.5
+        self.cosmo = FlatLambdaCDM(H0=70, Om0=0.3, Ob0=0.0)
+        self.lensCosmo = LensCosmo(
+            cosmo=self.cosmo, z_lens=self.z_lens, z_source=self.z_source
+        )
 
         # make class instances for a chosen lens model type
 
         # chose a lens model
-        lens_model_list = ["SPEP", "SHEAR"]
+        self.lens_model_list = ["SPEP", "SHEAR"]
         # make instance of LensModel class
-        lensModel = LensModel(
-            lens_model_list=lens_model_list,
-            cosmo=cosmo,
-            z_lens=z_lens,
-            z_source=z_source,
+        self.lensModel = LensModel(
+            lens_model_list=self.lens_model_list,
+            cosmo=self.cosmo,
+            z_lens=self.z_lens,
+            z_source=self.z_source,
         )
         # we require routines accessible in the LensModelExtensions class
         # make instance of LensEquationSolver to solve the lens equation
-        lensEquationSolver = LensEquationSolver(lensModel=lensModel)
+        self.lensEquationSolver = LensEquationSolver(lensModel=self.lensModel)
 
         # make choice of lens model
 
         # we chose a source position (in units angle)
-        x_source, y_source = 0.02, 0.01
-        x_source2, y_source2 = -0.02, 0.01
+        self.x_source, self.y_source = 0.02, 0.01
+        self.x_source2, self.y_source2 = -0.02, 0.01
         # we chose a lens model
-        kwargs_lens = [
+        self.kwargs_lens = [
             {
                 "theta_E": 1.0,
                 "e1": 0.1,
@@ -56,50 +55,35 @@ class TestImageLikelihood(object):
         # compute image positions and their (finite) magnifications
 
         # we solve for the image position(s) of the provided source position and lens model
-        x_img, y_img = lensEquationSolver.image_position_from_source(
-            kwargs_lens=kwargs_lens, sourcePos_x=x_source, sourcePos_y=y_source
+        self.x_img, self.y_img = self.lensEquationSolver.image_position_from_source(
+            kwargs_lens=self.kwargs_lens,
+            sourcePos_x=self.x_source,
+            sourcePos_y=self.y_source,
         )
-        x_img2, y_img2 = lensEquationSolver.image_position_from_source(
-            kwargs_lens=kwargs_lens, sourcePos_x=x_source2, sourcePos_y=y_source2
+        self.x_img2, self.y_img2 = self.lensEquationSolver.image_position_from_source(
+            kwargs_lens=self.kwargs_lens,
+            sourcePos_x=self.x_source2,
+            sourcePos_y=self.y_source2,
         )
+        self.point_source_list = ["LENSED_POSITION"]
+        self.kwargs_ps = [{"ra_image": self.x_img, "dec_image": self.y_img}]
+        self.pointSource = PointSource(point_source_type_list=self.point_source_list)
 
-        point_source_list = ["LENSED_POSITION"]
-        kwargs_ps = [{"ra_image": x_img, "dec_image": y_img}]
-        pointSource = PointSource(point_source_type_list=point_source_list)
-        t_days = lensModel.arrival_time(x_img, y_img, kwargs_lens)
+    def test_logL(self):
+        t_days = self.lensModel.arrival_time(self.x_img, self.y_img, self.kwargs_lens)
         time_delays_measured = t_days[1:] - t_days[0]
         time_delays_uncertainties = np.array([0.1, 0.1, 0.1])
-        self.td_likelihood = TimeDelayLikelihood(
+        td_likelihood = TimeDelayLikelihood(
             time_delays_measured,
             time_delays_uncertainties,
-            lens_model_class=lensModel,
-            point_source_class=pointSource,
-        )
-        kwargs_cosmo = {"D_dt": lensCosmo.ddt}
-        logL = self.td_likelihood.logL(
-            kwargs_lens=kwargs_lens, kwargs_ps=kwargs_ps, kwargs_cosmo=kwargs_cosmo
-        )
-        npt.assert_almost_equal(logL, 0, decimal=8)
-
-        # testing two sources
-        point_source_list2 = ["LENSED_POSITION", "LENSED_POSITION"]
-        kwargs_ps2 = [
-            {"ra_image": x_img, "dec_image": y_img},
-            {"ra_image": x_img2, "dec_image": y_img2},
-        ]
-        pointSource2 = PointSource(point_source_type_list=point_source_list2)
-        t2_days = lensModel.arrival_time(x_img2, y_img2, kwargs_lens)
-        time_delays_measured2 = t2_days[1:] - t2_days[0]
-        time_delays_uncertainties2 = np.array([0.1, 0.1, 0.1])
-        self.td_likelihood = TimeDelayLikelihood(
-            [time_delays_measured, time_delays_measured2],
-            [time_delays_uncertainties, time_delays_uncertainties2],
-            lens_model_class=lensModel,
-            point_source_class=pointSource2,
+            lens_model_class=self.lensModel,
+            point_source_class=self.pointSource,
         )
-        kwargs_cosmo = {"D_dt": lensCosmo.ddt}
-        logL = self.td_likelihood.logL(
-            kwargs_lens=kwargs_lens, kwargs_ps=kwargs_ps2, kwargs_cosmo=kwargs_cosmo
+        kwargs_cosmo = {"D_dt": self.lensCosmo.ddt}
+        logL = td_likelihood.logL(
+            kwargs_lens=self.kwargs_lens,
+            kwargs_ps=self.kwargs_ps,
+            kwargs_cosmo=kwargs_cosmo,
         )
         npt.assert_almost_equal(logL, 0, decimal=8)
 
@@ -108,12 +92,14 @@ class TestImageLikelihood(object):
         td_likelihood = TimeDelayLikelihood(
             time_delays_measured_new,
             time_delays_uncertainties,
-            lens_model_class=lensModel,
-            point_source_class=pointSource,
+            lens_model_class=self.lensModel,
+            point_source_class=self.pointSource,
         )
-        kwargs_cosmo = {"D_dt": lensCosmo.ddt}
+        kwargs_cosmo = {"D_dt": self.lensCosmo.ddt}
         logL = td_likelihood.logL(
-            kwargs_lens=kwargs_lens, kwargs_ps=kwargs_ps, kwargs_cosmo=kwargs_cosmo
+            kwargs_lens=self.kwargs_lens,
+            kwargs_ps=self.kwargs_ps,
+            kwargs_cosmo=kwargs_cosmo,
         )
         npt.assert_almost_equal(logL, -0.5, decimal=8)
 
@@ -122,11 +108,13 @@ class TestImageLikelihood(object):
         td_likelihood = TimeDelayLikelihood(
             time_delays_measured_new,
             time_delays_cov,
-            lens_model_class=lensModel,
-            point_source_class=pointSource,
+            lens_model_class=self.lensModel,
+            point_source_class=self.pointSource,
         )
         logL = td_likelihood.logL(
-            kwargs_lens=kwargs_lens, kwargs_ps=kwargs_ps, kwargs_cosmo=kwargs_cosmo
+            kwargs_lens=self.kwargs_lens,
+            kwargs_ps=self.kwargs_ps,
+            kwargs_cosmo=kwargs_cosmo,
         )
         npt.assert_almost_equal(logL, -0.5, decimal=8)
 
@@ -136,11 +124,13 @@ class TestImageLikelihood(object):
         td_likelihood = TimeDelayLikelihood(
             time_delays_measured_new,
             time_delays_uncertainties,
-            lens_model_class=lensModel,
-            point_source_class=pointSource,
+            lens_model_class=self.lensModel,
+            point_source_class=self.pointSource,
         )
         logL = td_likelihood.logL(
-            kwargs_lens=kwargs_lens, kwargs_ps=kwargs_ps, kwargs_cosmo=kwargs_cosmo
+            kwargs_lens=self.kwargs_lens,
+            kwargs_ps=self.kwargs_ps,
+            kwargs_cosmo=kwargs_cosmo,
         )
         npt.assert_almost_equal(logL, -(10**15), decimal=8)
 
@@ -149,6 +139,130 @@ class TestImageLikelihood(object):
 
         :return:
         """
+        # testing two sources
+        point_source_list2 = ["LENSED_POSITION", "LENSED_POSITION"]
+        kwargs_ps2 = [
+            {"ra_image": self.x_img, "dec_image": self.y_img},
+            {"ra_image": self.x_img2, "dec_image": self.y_img2},
+        ]
+
+        t_days = self.lensModel.arrival_time(self.x_img, self.y_img, self.kwargs_lens)
+        time_delays_measured = t_days[1:] - t_days[0]
+        time_delays_uncertainties = np.array([0.1, 0.1, 0.1])
+        pointSource2 = PointSource(point_source_type_list=point_source_list2)
+        t2_days = self.lensModel.arrival_time(
+            self.x_img2, self.y_img2, self.kwargs_lens
+        )
+        time_delays_measured2 = t2_days[1:] - t2_days[0]
+        time_delays_uncertainties2 = np.array([0.1, 0.1, 0.1])
+        td_likelihood = TimeDelayLikelihood(
+            [time_delays_measured, time_delays_measured2],
+            [time_delays_uncertainties, time_delays_uncertainties2],
+            lens_model_class=self.lensModel,
+            point_source_class=pointSource2,
+        )
+        kwargs_cosmo = {"D_dt": self.lensCosmo.ddt}
+        logL = td_likelihood.logL(
+            kwargs_lens=self.kwargs_lens,
+            kwargs_ps=kwargs_ps2,
+            kwargs_cosmo=kwargs_cosmo,
+        )
+        npt.assert_almost_equal(logL, 0, decimal=8)
+
+        time_delays_measured_new = copy.deepcopy(time_delays_measured)
+        time_delays_measured_new[0] += 0.1
+        td_likelihood = TimeDelayLikelihood(
+            time_delays_measured_new,
+            time_delays_uncertainties,
+            lens_model_class=self.lensModel,
+            point_source_class=self.pointSource,
+        )
+        kwargs_cosmo = {"D_dt": self.lensCosmo.ddt}
+        logL = td_likelihood.logL(
+            kwargs_lens=self.kwargs_lens,
+            kwargs_ps=kwargs_ps2,
+            kwargs_cosmo=kwargs_cosmo,
+        )
+        npt.assert_almost_equal(logL, -0.5, decimal=8)
+
+        # Test if the time delay measurement bool list is used correctly
+        td_likelihood = TimeDelayLikelihood(
+            [time_delays_measured, time_delays_measured2],
+            [time_delays_uncertainties, time_delays_uncertainties2],
+            lens_model_class=self.lensModel,
+            point_source_class=pointSource2,
+            time_delay_measurement_bool_list=[[True, True, True], [True, False, False]],
+        )
+        kwargs_cosmo = {"D_dt": self.lensCosmo.ddt}
+        logL = td_likelihood.logL(
+            kwargs_lens=self.kwargs_lens,
+            kwargs_ps=kwargs_ps2,
+            kwargs_cosmo=kwargs_cosmo,
+        )
+        npt.assert_almost_equal(logL, 0, decimal=8)
+
+        # Test if the time delay measurement bool list is used correctly if only one bool per point source is given
+        td_likelihood = TimeDelayLikelihood(
+            [time_delays_measured, time_delays_measured2],
+            [time_delays_uncertainties, time_delays_uncertainties2],
+            lens_model_class=self.lensModel,
+            point_source_class=pointSource2,
+            time_delay_measurement_bool_list=[True, True],
+        )
+        kwargs_cosmo = {"D_dt": self.lensCosmo.ddt}
+        logL = td_likelihood.logL(
+            kwargs_lens=self.kwargs_lens,
+            kwargs_ps=kwargs_ps2,
+            kwargs_cosmo=kwargs_cosmo,
+        )
+        npt.assert_almost_equal(logL, 0, decimal=8)
+
+        # Test if the time delay measurement bool list is used correctly with a covariance matrix
+        time_delays_cov = np.diag([0.1, 0.1, 0.1]) ** 2
+        time_delays_cov2 = np.diag([0.1, 0.1, np.nan]) ** 2
+        time_delays_measured2[2] = np.nan
+        td_likelihood = TimeDelayLikelihood(
+            [time_delays_measured, time_delays_measured2],
+            [time_delays_cov, time_delays_cov2],
+            lens_model_class=self.lensModel,
+            point_source_class=pointSource2,
+            time_delay_measurement_bool_list=[[True, True, True], [True, True, False]],
+        )
+        kwargs_cosmo = {"D_dt": self.lensCosmo.ddt}
+        logL = td_likelihood.logL(
+            kwargs_lens=self.kwargs_lens,
+            kwargs_ps=kwargs_ps2,
+            kwargs_cosmo=kwargs_cosmo,
+        )
+        npt.assert_almost_equal(logL, 0, decimal=8)
+
+        # Test the raise statement
+        with pytest.raises(ValueError):
+            td_likelihood = TimeDelayLikelihood(
+                [time_delays_measured, time_delays_measured2],
+                [time_delays_uncertainties, time_delays_uncertainties2],
+                lens_model_class=self.lensModel,
+                point_source_class=pointSource2,
+                time_delay_measurement_bool_list=[[True, True, True], [True, False]],
+            )
+
+        with pytest.raises(ValueError):
+            td_likelihood = TimeDelayLikelihood(
+                [time_delays_measured, time_delays_measured2],
+                [time_delays_uncertainties, time_delays_uncertainties2],
+                lens_model_class=self.lensModel,
+                point_source_class=pointSource2,
+                time_delay_measurement_bool_list=[[True, True, True]],
+            )
+
+        with pytest.raises(ValueError):
+            td_likelihood = TimeDelayLikelihood(
+                [time_delays_measured, time_delays_measured2],
+                [time_delays_uncertainties, time_delays_uncertainties2],
+                lens_model_class=self.lensModel,
+                point_source_class=pointSource2,
+                time_delay_measurement_bool_list=["aaa", "bbb"],
+            )
 
 
 if __name__ == "__main__":