Work on cluster lensing, add LOS lens draw method to lens_pop.py, optimize _image_position_from_source for multi_plane="Source" lenses

slsim/Deflectors/DeflectorPopulation/cluster_deflectors.py

@@ -141,7 +141,8 @@ def draw_deflector(self, index=None):
         :type cored: True for cored, False for cuspy profile
         :return: dictionary of complete parameterization of deflector
         """
-        index = random.randint(0, self._num_select - 1)
+        if index is None:
+            index = random.randint(0, self._num_select - 1)
         deflector = self.draw_cluster(index)
         members = self.draw_members(deflector["cluster_id"], **self.kwargs_draw_members)
         deflector["subhalos"] = members

slsim/Deflectors/deflector.py

@@ -10,6 +10,7 @@
 from lenstronomy.Cosmo.lens_cosmo import LensCosmo
 from lenstronomy.Analysis.lens_profile import LensProfileAnalysis
 from lenstronomy.LensModel.lens_model import LensModel
+from lenstronomy.LensModel.lens_model_extensions import LensModelExtensions
 
 _SUPPORTED_DEFLECTORS = ["EPL", "EPL_SERSIC", "NFW_HERNQUIST", "NFW_CLUSTER"]
 JAX_PROFILES = [
@@ -288,3 +289,44 @@ def theta_e_infinity(self, cosmo, multi_plane=None, use_jax=True):
             theta_E_infinity = np.nan_to_num(theta_E_infinity, nan=0)
         self._theta_e_infinity = theta_E_infinity
         return theta_E_infinity
+
+    def critical_curves_caustics_list(
+        self, z_source, cosmo, kwargs_critical_curve_caustics=None
+    ):
+        """Returns list of critical curves and caustics for a source at
+        `z_source`
+
+        :param z_source: redshift at which to compute curves
+        :param cosmo: astropy.cosmology instance
+        :param kwargs_critical_curve_caustics: arguments passed into the `critical_curve_caustics` function from LensModelExtensions. Keys:
+        - compute_window: window size in arcsec where the critical curve is computed
+        - grid_scale: numerical grid spacing of the computation of the critical curves
+        - center_x: float, center of the window to compute critical curves and caustics
+        - center_y: float, center of the window to compute critical curves and caustics
+        - kwargs_lens: lens model kwargs
+        """
+
+        lens_cosmo = LensCosmo(
+            z_lens=self.redshift,
+            z_source=z_source,
+            cosmo=cosmo,
+        )
+
+        lens_mass_model_list, model_params = self.mass_model_lenstronomy(lens_cosmo)
+
+        lens_model = LensModel(
+            lens_model_list=lens_mass_model_list,
+            cosmo=cosmo,
+            z_lens=self.redshift,
+            z_source=z_source,
+            multi_plane=False,
+        )
+
+        lens_model_ext = LensModelExtensions(lens_model)
+        ra_crit_list, dec_crit_list, ra_caustic_list, dec_caustic_list = (
+            lens_model_ext.critical_curve_caustics(
+                model_params, **kwargs_critical_curve_caustics
+            )
+        )
+
+        return ra_crit_list, dec_crit_list, ra_caustic_list, dec_caustic_list

slsim/Lenses/lens.py

@@ -43,6 +43,7 @@ def __init__(
         shear=True,
         convergence=True,
         field_galaxies=None,
+        create_field_galaxies=False,
     ):
         """
 
@@ -75,8 +76,21 @@ def __init__(
             Instances should be generated via :meth:`slsim.Sources.SourcePopulation.Galaxies.draw_field_galaxies`.
             If None, no field galaxies are included.
         :type field_galaxies: list[`slsim.Sources.source.Source`] or None
-
+        :param create_field_galaxies: If True, make all source galaxies in front of `deflector_class.redshift` into field galaxies
+        :type create_field_galaxies: bool
         """
+
+        if create_field_galaxies and isinstance(source_class, list):
+            if field_galaxies is None:
+                field_galaxies = []
+
+            field_galaxies += [
+                s for s in source_class if s.redshift < deflector_class.redshift
+            ]
+            source_class = [
+                s for s in source_class if s.redshift > deflector_class.redshift
+            ]
+
         LensedSystemBase.__init__(
             self,
             source_class=source_class,
@@ -194,7 +208,12 @@ def _image_position_from_source(self, x_source, y_source, source_index):
             list], [DEC list]
         """
         lens_model_class, kwargs_lens = self.deflector_mass_model_lenstronomy(
-            source_index=source_index
+            source_index=source_index,
+            multi_plane=self.multi_plane
+            in [
+                "Deflector",
+                "Both",
+            ],  # unnecessary for source, makes calculations much faster
         )
         lens_eq_solver = LensEquationSolver(lens_model_class)
         point_source_pos_x, point_source_pos_y = x_source, y_source
@@ -1410,7 +1429,7 @@ def lenstronomy_kwargs(
 
         return kwargs_model, kwargs_params
 
-    def deflector_mass_model_lenstronomy(self, source_index=None):
+    def deflector_mass_model_lenstronomy(self, source_index=None, multi_plane=None):
         """Returns lens model instance and parameters in lenstronomy
         conventions.
 
@@ -1493,7 +1512,7 @@ def deflector_mass_model_lenstronomy(self, source_index=None):
             z_source=z_source,
             z_source_convention=self.max_redshift_source_class.redshift,
             use_jax=use_jax,
-            multi_plane=bool(self.multi_plane),
+            multi_plane=bool(self.multi_plane) if multi_plane is None else multi_plane,
         )
 
         return lens_model, self._kwargs_lens

slsim/Lenses/lens_pop.py

@@ -8,6 +8,8 @@
 from slsim.LOS.los_pop import LOSPop
 from slsim.Deflectors.DeflectorPopulation.deflectors_base import DeflectorsBase
 from slsim.Lenses.lensed_population_base import LensedPopulationBase
+from lenstronomy.LensModel.lens_model_extensions import LensModelExtensions
+from matplotlib.path import Path
 
 from tqdm import tqdm
 
@@ -107,6 +109,135 @@ def select_lens_at_random(self, test_area=None, verbose=False, **kwargs_lens_cut
                 return gg_lens
             n += 1
 
+    def select_lens_at_random_multi_source(
+        self,
+        source_area,
+        verbose=False,
+        min_num_sources=1,
+        return_only_multiply_imaged_sources=False,
+        **kwargs_lens_cut
+    ):
+        """Draw a random lens with an entire source field spanning
+        `source_area`, with at least min_num_sources satisfying
+        kwargs_lens_cut.
+
+        :param sky_area: Sky area to draw sources from
+        :type sky_area: (astropy.units.Quantity)
+        :param kwargs_lens_cut: Dictionary of cuts that one wants to apply to the lens.
+                                Example: kwargs_lens_cut = {
+                                "min_image_separation": 0.5,
+                                "max_image_separation": 10,
+                                "mag_arc_limit": {"i": 24},
+                                "second_brightest_image_cut": {"i": 24}}.
+                                All these cuts are optional.
+        :type kwargs_lens_cut: dict
+        :param min_num_sources: Minimum number of lensed sources that must satisfy kwargs_lens_cut
+        :type min_num_sources: int
+        :param return_only_multiply_imaged_sources: Return a lens class only with sources that are multiply imaged. Default False.
+        :type return_only_multiply_imaged_sources: bool
+        :param verbose: print statements added
+        :type verbose: bool
+        :return: Lens() instance with parameters of the deflector and lens and source field.
+        :rtype: Lens
+        """
+
+        # utility function - returns True if point is inside one of the paths in the ra and dec lists, or within 3 arcseconds outside. Used to cut sources for image checking.
+        def _in_caustic_or_close_outside(ra_caustic_list, dec_caustic_list, s):
+            point = s.extended_source_position
+
+            for ra, dec in zip(ra_caustic_list, dec_caustic_list):
+                curve_path = Path(np.column_stack((ra, dec)))
+
+                area = (np.max(ra) - np.min(ra)) * (np.max(dec) - np.min(dec))
+                if area < 1:
+                    continue
+
+                if curve_path.contains_point(point):
+                    return True
+
+                if np.min((ra - point[0]) ** 2 + (dec - point[1]) ** 2) < 3**2:
+                    return True
+
+        n = 0
+        while True:
+            # draw random deflector
+            _deflector = self._lens_galaxies.draw_deflector()
+
+            ### compute caustics at high redshift to filter source galaxies for validity checking
+            _, _, ra_caustic_list, dec_caustic_list = (
+                _deflector.critical_curves_caustics_list(
+                    10,
+                    self.cosmo,
+                    {
+                        "compute_window": np.sqrt(
+                            source_area.to_value("arcsec2") / np.pi
+                        )
+                        * 2,
+                        "grid_scale": 0.5,
+                    },
+                )
+            )
+
+            if len(ra_caustic_list) == 0:
+                continue
+
+            # draw all sources, and filter onces near caustics for validity checking
+            _source = self._sources.draw_galaxies(source_area)
+            _source_cut = [
+                s
+                for s in _source
+                if s.redshift > _deflector.redshift
+                and _in_caustic_or_close_outside(ra_caustic_list, dec_caustic_list, s)
+            ]
+
+            if len(_source_cut) < min_num_sources:
+                continue
+
+            # lens only with sources near caustics to speed up validity checking
+            test_lens = Lens(
+                deflector_class=_deflector,
+                source_class=_source_cut,
+                cosmo=self.cosmo,
+                use_jax=self._use_jax,
+                multi_plane="Source",
+                create_field_galaxies=True,
+            )
+
+            test_res = test_lens.validity_test(**kwargs_lens_cut)
+            if not isinstance(test_res, dict):
+                test_res = {0: test_res}
+
+            if len([x for x in test_res.values() if x]) >= min_num_sources:
+                if verbose is True:
+                    print("selected lens after %s tries." % n)
+
+                if not return_only_multiply_imaged_sources:
+                    # final lens with all sources
+                    return Lens(
+                        deflector_class=_deflector,
+                        source_class=_source,
+                        cosmo=self.cosmo,
+                        use_jax=self._use_jax,
+                        multi_plane="Source",
+                        create_field_galaxies=True,
+                    )
+                else:
+                    # only sources that are multiply imaged
+                    return Lens(
+                        deflector_class=_deflector,
+                        source_class=[
+                            _source_cut[i]
+                            for i in range(len(_source_cut))
+                            if test_res[i]
+                        ],
+                        cosmo=self.cosmo,
+                        use_jax=self._use_jax,
+                        multi_plane="Source",
+                        create_field_galaxies=True,
+                    )
+
+            n += 1
+
     def _draw_source(self, mag_arc_limit=None, magnification_limit=2, **kwargs):
         """Draw from source population considering some additional constraints
         to be fulfilled.

slsim/Pipelines/skypy_pipeline.py

@@ -2,6 +2,7 @@
 from skypy.pipeline import Pipeline
 import tempfile
 import slsim.Util.param_util as util
+from astropy.table import vstack
 
 
 class SkyPyPipeline:
@@ -111,3 +112,12 @@ def red_galaxies(self):
         :rtype: list of dict
         """
         return self._pipeline["red"]
+
+    @property
+    def all_galaxies(self):
+        """Skypy pipeline for red and blue galaxies.
+
+        :return: list of red and blue galaxies
+        :rtype: list of dict
+        """
+        return vstack([self.red_galaxies, self.blue_galaxies])

tests/test_Lenses/test_lens_pop.py

@@ -226,6 +226,63 @@ def test_cluster_lens_pop_instance():
     assert pes_lens_class.deflector_velocity_dispersion() > 250
 
 
+def test_cluster_lens_pop_instance_multi_source():
+    np.random.seed(41)
+    cosmo = FlatLambdaCDM(H0=70, Om0=0.3)
+    sky_area = Quantity(value=100**2, unit="arcsec2")
+
+    kwargs_deflector_cut = {"z_min": 0.2, "z_max": 1.0}
+    kwargs_source_cut = {"band": "g", "band_max": 28, "z_min": 0.25, "z_max": 5.0}
+
+    path = os.path.dirname(__file__)
+    module_path = os.path.dirname(os.path.dirname(path))
+    cluster_catalog_path = os.path.join(
+        module_path, "data/redMaPPer/clusters_example.fits"
+    )
+    members_catalog_path = os.path.join(
+        module_path, "data/redMaPPer/members_example.fits"
+    )
+    cluster_catalog = Table.read(cluster_catalog_path)
+    members_catalog = Table.read(members_catalog_path)
+
+    lens_clusters = deflectors.ClusterDeflectors(
+        cluster_list=cluster_catalog,
+        members_list=members_catalog,
+        galaxy_list=galaxy_simulation_pipeline.all_galaxies,
+        kwargs_cut=kwargs_deflector_cut,
+        kwargs_mass2light={},
+        cosmo=cosmo,
+        sky_area=sky_area,
+        kwargs_draw_members={"max_dist": 350},
+    )
+
+    source_galaxies = sources.Galaxies(
+        galaxy_list=galaxy_simulation_pipeline.all_galaxies,
+        kwargs_cut=kwargs_source_cut,
+        cosmo=cosmo,
+        sky_area=sky_area,
+        catalog_type="skypy",
+    )
+
+    cluster_lens_pop = LensPop(
+        deflector_population=lens_clusters,
+        source_population=source_galaxies,
+        cosmo=cosmo,
+        sky_area=sky_area,
+        use_jax=use_jax,
+    )
+
+    kwargs_lens_cut = {"min_image_separation": 1.0, "max_image_separation": 100.0}
+    pes_lens_class = cluster_lens_pop.select_lens_at_random_multi_source(
+        source_area=sky_area, min_num_sources=2, **kwargs_lens_cut
+    )
+    assert pes_lens_class.deflector.deflector_type == "NFW_CLUSTER"
+    kwargs_model, kwargs_params = pes_lens_class.lenstronomy_kwargs(band="g")
+    assert len(kwargs_model["lens_model_list"]) >= 3  # halo, 1>= subhalo, LoS
+    assert len(kwargs_model["lens_light_model_list"]) >= 1  # 1>= member galaxy
+    assert pes_lens_class.deflector_velocity_dispersion() > 250
+
+
 def test_galaxies_lens_pop_instance():
     cosmo = FlatLambdaCDM(H0=70, Om0=0.3)
     sky_area = Quantity(value=0.001, unit="deg2")