update to positions_likelihood_list

Author: Jammy2211

autolens/analysis/analysis/dataset.py

@@ -1,9 +1,9 @@
 import os
 import logging
-from typing import Optional, Union
+from typing import List, Optional, Union
 
 from autoconf import conf
-from autoconf.dictable import to_dict, output_to_json
+from autoconf.dictable import output_to_json
 
 import autofit as af
 import autoarray as aa
@@ -27,8 +27,8 @@ class AnalysisDataset(AgAnalysisDataset, AnalysisLens):
     def __init__(
         self,
         dataset,
-        positions_likelihood: Optional[
-            Union[PositionsLHResample, PositionsLHPenalty]
+        positions_likelihood_list: Optional[
+            List[Union[PositionsLHResample, PositionsLHPenalty]]
         ] = None,
         adapt_image_maker: Optional[ag.AdaptImageMaker] = None,
         cosmology: ag.cosmo.LensingCosmology = ag.cosmo.Planck15(),
@@ -49,10 +49,10 @@ def __init__(
         ----------
         dataset
             The imaging, interferometer or other dataset that the model if fitted too.
-        positions_likelihood
-            An object which alters the likelihood function to include a term which accounts for whether
-            image-pixel coordinates in arc-seconds corresponding to the multiple images of the lensed source galaxy
-            trace close to one another in the source-plane.
+        positions_likelihood_list
+            Alters the likelihood function to include a term which accounts for whether image-pixel coordinates in
+            arc-seconds corresponding to the multiple images of each lensed source galaxy trace close to one another in
+            their source-plane.
         adapt_images
             Contains the adapt-images which are used to make a pixelization's mesh and regularization adapt to the
             reconstructed galaxy's morphology.
@@ -78,7 +78,7 @@ def __init__(
 
         AnalysisLens.__init__(
             self=self,
-            positions_likelihood=positions_likelihood,
+            positions_likelihood_list=positions_likelihood_list,
             cosmology=cosmology,
         )
 

autolens/analysis/analysis/lens.py

@@ -1,6 +1,6 @@
 import logging
 import numpy as np
-from typing import Dict, Optional, Union
+from typing import Dict, List, Optional, Union
 
 import autofit as af
 import autoarray as aa
@@ -22,8 +22,8 @@
 class AnalysisLens:
     def __init__(
         self,
-        positions_likelihood: Optional[
-            Union[PositionsLHResample, PositionsLHPenalty]
+        positions_likelihood_list: Optional[
+            List[Union[PositionsLHResample, PositionsLHPenalty]]
         ] = None,
         cosmology: ag.cosmo.LensingCosmology = ag.cosmo.Planck15(),
     ):
@@ -43,7 +43,7 @@ def __init__(
             The Cosmology assumed for this analysis.
         """
         self.cosmology = cosmology
-        self.positions_likelihood = positions_likelihood
+        self.positions_likelihood_list = positions_likelihood_list
 
     def tracer_via_instance_from(
         self,
@@ -123,10 +123,13 @@ def log_likelihood_positions_overwrite_from(
         The penalty value of the positions log likelihood, if the positions do not trace close in the source plane,
         else a None is returned to indicate there is no penalty.
         """
-        if self.positions_likelihood is not None:
+        if self.positions_likelihood_list is not None:
             try:
-                return self.positions_likelihood.log_likelihood_function_positions_overwrite(
-                    instance=instance, analysis=self
-                )
+                log_likelihood = 0.0
+                for positions_likelihood in self.positions_likelihood_list:
+                    log_likelihood += positions_likelihood.log_likelihood_function_positions_overwrite(
+                        instance=instance, analysis=self
+                    )
+                return log_likelihood
             except (ValueError, np.linalg.LinAlgError) as e:
                 raise exc.FitException from e

autolens/analysis/positions.py

@@ -26,7 +26,7 @@ def __init__(
         self,
         threshold: float,
         positions: Optional[aa.Grid2DIrregular] = None,
-        plane_redshift_positions_dict: Optional[Dict[float, aa.Grid2DIrregular]] = None,
+        plane_redshift: Optional[float] = None,
     ):
         """
         The `PositionsLH` objects add a penalty term to the likelihood of the **PyAutoLens** `log_likelihood_function`
@@ -44,9 +44,8 @@ def __init__(
         The default behaviour assumes a single lens plane and single source plane, meaning the input `positions`
         are the image-plane coordinates of one source galaxy at a specifc redshift.
 
-        For multiple source planes, the `plane_redshift_positions_dict` dictionary can be used to input multiple
-        sets of image positions, where the key is the redshift of the source plane and the value is the
-        corresponding image-plane coordinates.
+        For multiple source planes, the `plane_redshift` can be used to pair the image-plane positions with the
+        redshift of the source plane they trace too.
 
         Parameters
         ----------
@@ -56,31 +55,21 @@ def __init__(
         threshold
             If the maximum separation of any two source plane coordinates is above the threshold the penalty term
             is applied.
-        plane_redshift_positions_dict
-            A dictionary of plane redshifts and the corresponding image-plane coordinates of the lensed source
-            multiple images. The key is the redshift of the source plane and the value is the corresponding
-            image-plane coordinates.
+        plane_redshift
+            The plane redshift of the lensed source multiple images, which is only required if position threshold
+            for a double source plane lens system is being used where the specific plane is required.
         """
 
         self.positions = positions
         self.threshold = threshold
+        self.plane_redshift = plane_redshift
 
-        self.plane_redshift_positions_dict = plane_redshift_positions_dict
-
-        if plane_redshift_positions_dict is None:
-            self.plane_redshift_positions_dict = {None: positions}
-
-        if len(self.plane_redshift_positions_dict) == 1:
-            self.positions = list(self.plane_redshift_positions_dict.values())[0]
-
-        for positions in self.plane_redshift_positions_dict.values():
-
-            if len(positions) == 1:
-                raise exc.PositionsException(
-                    f"The positions input into the PositionsLikelihood object have length one "
-                    f"(e.g. it is only one (y,x) coordinate and therefore cannot be compared with other images).\n\n"
-                    "Please input more positions into the Positions."
-                )
+        if len(positions) == 1:
+            raise exc.PositionsException(
+                f"The positions input into the PositionsLikelihood object have length one "
+                f"(e.g. it is only one (y,x) coordinate and therefore cannot be compared with other images).\n\n"
+                "Please input more positions into the Positions."
+            )
 
     def log_likelihood_function_positions_overwrite(
         self, instance: af.ModelInstance, analysis: AnalysisDataset
@@ -107,29 +96,22 @@ def output_positions_info(self, output_path: str, tracer: Tracer):
         """
         with open_(path.join(output_path, "positions.info"), "w+") as f:
 
-            plane_index = 0
-
-            for plane_redshift, positions in self.plane_redshift_positions_dict.items():
-
-                positions_fit = SourceMaxSeparation(
-                    data=positions, noise_map=None, tracer=tracer, plane_redshift=plane_redshift
-                )
-
-                distances = positions_fit.data.distances_to_coordinate_from(
-                    coordinate=(0.0, 0.0)
-                )
+            positions_fit = SourceMaxSeparation(
+                data=self.positions, noise_map=None, tracer=tracer, plane_redshift=self.plane_redshift
+            )
 
-                f.write(f"Plane Index: {plane_index} \n")
-                f.write(f"Plane Redshift: {plane_redshift} \n")
-                f.write(f"Positions: \n {positions} \n\n")
-                f.write(f"Radial Distance from (0.0, 0.0): \n {distances} \n\n")
-                f.write(f"Threshold = {self.threshold} \n")
-                f.write(
-                    f"Max Source Plane Separation of Maximum Likelihood Model = {positions_fit.max_separation_of_plane_positions}"
-                )
-                f.write("")
+            distances = positions_fit.data.distances_to_coordinate_from(
+                coordinate=(0.0, 0.0)
+            )
 
-                plane_index += 1
+            f.write(f"Plane Redshift: {self.plane_redshift} \n")
+            f.write(f"Positions: \n {self.positions} \n\n")
+            f.write(f"Radial Distance from (0.0, 0.0): \n {distances} \n\n")
+            f.write(f"Threshold = {self.threshold} \n")
+            f.write(
+                f"Max Source Plane Separation of Maximum Likelihood Model = {positions_fit.max_separation_of_plane_positions}"
+            )
+            f.write("")
 
 
 class PositionsLHResample(AbstractPositionsLH):
@@ -179,20 +161,18 @@ def log_likelihood_function_positions_overwrite(
         if not tracer.has(cls=ag.mp.MassProfile) or len(tracer.planes) == 1:
             return
 
-        for plane_redshift, positions in self.plane_redshift_positions_dict.items():
-
-            positions_fit = SourceMaxSeparation(
-                data=positions,
-                noise_map=None,
-                tracer=tracer,
-                plane_redshift=plane_redshift,
-            )
+        positions_fit = SourceMaxSeparation(
+            data=self.positions,
+            noise_map=None,
+            tracer=tracer,
+            plane_redshift=self.plane_redshift,
+        )
 
-            if not positions_fit.max_separation_within_threshold(self.threshold):
-                if os.environ.get("PYAUTOFIT_TEST_MODE") == "1":
-                    return
+        if not positions_fit.max_separation_within_threshold(self.threshold):
+            if os.environ.get("PYAUTOFIT_TEST_MODE") == "1":
+                return
 
-                raise exc.RayTracingException
+            raise exc.RayTracingException
 
 
 class PositionsLHPenalty(AbstractPositionsLH):
@@ -201,7 +181,7 @@ def __init__(
         threshold: float,
         log_likelihood_penalty_factor: float = 1e8,
         positions: Optional[aa.Grid2DIrregular] = None,
-        plane_redshift_positions_dict: Optional[Dict[int, aa.Grid2DIrregular]] = None,
+        plane_redshift: Optional[float] = None,
     ):
         """
         The `PositionsLH` objects add a penalty term to the likelihood of the **PyAutoLens** `log_likelihood_function`
@@ -231,15 +211,14 @@ def __init__(
         log_likelihood_penalty_factor
             A factor which multiplies how far source pixels do not trace within the threshold of one another, with a
             larger factor producing a larger penalty making the non-linear parameter space gradient steeper.
-        plane_redshift_positions_dict
-            A dictionary of plane redshifts and the corresponding image-plane coordinates of the lensed source
-            multiple images. The key is the redshift of the source plane and the value is the corresponding
-            image-plane coordinates.
+        plane_redshift
+            The plane redshift of the lensed source multiple images, which is only required if position threshold
+            for a double source plane lens system is being used where the specific plane is required.
         """
         super().__init__(
             positions=positions,
             threshold=threshold,
-            plane_redshift_positions_dict=plane_redshift_positions_dict,
+            plane_redshift=plane_redshift
         )
 
         self.log_likelihood_penalty_factor = log_likelihood_penalty_factor
@@ -318,21 +297,19 @@ def log_likelihood_penalty_from(self, tracer: Tracer) -> Optional[float]:
 
         log_likelihood_penalty = 0.0
 
-        for plane_redshift, positions in self.plane_redshift_positions_dict.items():
-
-            positions_fit = SourceMaxSeparation(
-                data=positions,
-                noise_map=None,
-                tracer=tracer,
-                plane_redshift=plane_redshift,
-            )
+        positions_fit = SourceMaxSeparation(
+            data=self.positions,
+            noise_map=None,
+            tracer=tracer,
+            plane_redshift=self.plane_redshift,
+        )
 
-            if not positions_fit.max_separation_within_threshold(self.threshold):
+        if not positions_fit.max_separation_within_threshold(self.threshold):
 
-                log_likelihood_penalty += self.log_likelihood_penalty_factor * (
-                    positions_fit.max_separation_of_plane_positions
-                    - self.threshold
-                )
+            log_likelihood_penalty += self.log_likelihood_penalty_factor * (
+                positions_fit.max_separation_of_plane_positions
+                - self.threshold
+            )
 
         return log_likelihood_penalty
 

test_autolens/imaging/model/test_analysis_imaging.py

@@ -57,7 +57,7 @@ def test__positions__resample__raises_exception(masked_imaging_7x7):
     )
 
     analysis = al.AnalysisImaging(
-        dataset=masked_imaging_7x7, positions_likelihood=positions_likelihood
+        dataset=masked_imaging_7x7, positions_likelihood_list=[positions_likelihood]
     )
 
     instance = model.instance_from_unit_vector([])
@@ -89,7 +89,7 @@ def test__positions__likelihood_overwrites__changes_likelihood(masked_imaging_7x
     )
 
     analysis = al.AnalysisImaging(
-        dataset=masked_imaging_7x7, positions_likelihood=positions_likelihood
+        dataset=masked_imaging_7x7, positions_likelihood_list=[positions_likelihood]
     )
     analysis_log_likelihood = analysis.log_likelihood_function(instance=instance)
 
@@ -116,17 +116,15 @@ def test__positions__likelihood_overwrites__changes_likelihood__double_source_pl
 
     instance = model.instance_from_unit_vector([])
 
-    plane_redshift_positions_dict = {
-        1.0: al.Grid2DIrregular([(1.0, 100.0), (200.0, 2.0)]),
-        2.0: al.Grid2DIrregular([(1.0, 100.0), (200.0, 2.0)])
-    }
-
-    positions_likelihood = al.PositionsLHPenalty(
-        plane_redshift_positions_dict=plane_redshift_positions_dict, threshold=0.01
+    positions_likelihood_0 = al.PositionsLHPenalty(
+        plane_redshift=1.0, positions=al.Grid2DIrregular([(1.0, 100.0), (200.0, 2.0)]), threshold=0.01
+    )
+    positions_likelihood_1 = al.PositionsLHPenalty(
+        plane_redshift=2.0, positions=al.Grid2DIrregular([(1.0, 100.0), (200.0, 2.0)]), threshold=0.01
     )
 
     analysis = al.AnalysisImaging(
-        dataset=masked_imaging_7x7, positions_likelihood=positions_likelihood
+        dataset=masked_imaging_7x7, positions_likelihood_list=[positions_likelihood_0, positions_likelihood_1]
     )
     analysis_log_likelihood = analysis.log_likelihood_function(instance=instance)