ler 0.4.1__py3-none-any.whl → 0.4.3__py3-none-any.whl

ler/lens_galaxy_population/lens_galaxy_parameter_distribution.py

@@ -2,187 +2,185 @@
 """
 This module contains the LensGalaxyPopulation class, which is used to sample lens galaxy parameters, source parameters conditioned on the source being strongly lensed. \n
 The class inherits from the ImageProperties class, which is used calculate image properties (magnification, timedelays, source position, image position, morse phase). \n
-Either the class takes in initialized CompactBinaryPopulation class as input or inherits the CompactBinaryPopulation class with default params (if no input) \n
+Either the class takes in initialized CBCSourceParameterDistribution class as input or inherits the CBCSourceParameterDistribution class with default params (if no input) \n
 """
 
 import warnings
 warnings.filterwarnings("ignore")
 import numpy as np
-from numba import njit
 from scipy.integrate import quad
-from lenstronomy.Util.param_util import phi_q2_ellipticity
-
-# for redshift to luminosity distance conversion
-from astropy.cosmology import LambdaCDM
+# from lenstronomy.Util.param_util import phi_q2_ellipticity
+from multiprocessing import Pool
+from tqdm import tqdm
 
 # the following .py file will be called if they are not given in the class initialization
 from ..gw_source_population import CBCSourceParameterDistribution
 from .optical_depth import OpticalDepth
 from ..image_properties import ImageProperties
-from ..utils import add_dictionaries_together, trim_dictionary
-from .jit_functions import phi_cut_SIE, velocity_dispersion_z_dependent, lens_redshift_SDSS_catalogue
-
+from ..utils import add_dictionaries_together, trim_dictionary, FunctionConditioning, interpolator_pickle_path
+from .jit_functions import phi_cut_SIE, phi_q2_ellipticity_hemanta
+from .mp import cross_section_mp, rjs_sie_mp
 
 class LensGalaxyParameterDistribution(CBCSourceParameterDistribution, ImageProperties, OpticalDepth):
     """
-    Class to sample lens galaxy parameters, source parameters conditioned on the source being strongly lensed, and image properties \n
-
-    Parameters
-    ----------
-    npool : `int`
-        number of processors to use
-    z_min : `float`
-        minimum redshift
-    z_max : `float`
-        maximum redshift    
-    cosmology : `astropy.cosmology`
-        Cosmology to use
-        default: None/astropy.cosmology.FlatLambdaCDM(H0=70, Om0=0.3)
-    event_type : `str`
-        Type of event to generate.
-        e.g. 'BBH', 'BNS', 'NSBH'
-        default: 'BBH'
-    lens_type : `str`
-        Type of lens galaxy to generate.
-        default: 'epl_galaxy'
-    lens_functions, lens_priors, lens_priors_params : `dict`, `dict`, `dict`
-        dictionary of lens functions, priors, and priors parameters
-        Check for default/available lens functions, priors and corresponding input parameters by running,\n
+        Class to sample lens galaxy parameters, source parameters conditioned on the source being strongly lensed, and image properties \n
+
+        Parameters
+        ----------
+        npool : `int`
+            number of processors to use
+        z_min : `float`
+            minimum redshift
+        z_max : `float`
+            maximum redshift    
+        cosmology : `astropy.cosmology`
+            Cosmology to use
+            default: None/astropy.cosmology.FlatLambdaCDM(H0=70, Om0=0.3)
+        event_type : `str`
+            Type of event to generate.
+            e.g. 'BBH', 'BNS', 'NSBH'
+            default: 'BBH'
+        lens_type : `str`
+            Type of lens galaxy to generate.
+            default: 'epl_shear_galaxy'
+        lens_functions, lens_priors, lens_priors_params : `dict`, `dict`, `dict`
+            dictionary of lens functions, priors, and priors parameters
+            Check for default/available lens functions, priors and corresponding input parameters by running,\n
+            >>> from ler.lens_galaxy_population import LensGalaxyParameterDistribution
+            >>> lens = LensGalaxyParameterDistribution()
+            >>> print(lens.lens_functions)
+            >>> print(lens.lens_priors)
+            >>> print(lens.lens_priors_params)
+        directory : `str`
+            directory to store the interpolators
+            default: './interpolator_pickle'
+        **kwargs : 
+            keyword arguments to pass to the parent classes
+
+        Examples
+        --------
         >>> from ler.lens_galaxy_population import LensGalaxyParameterDistribution
         >>> lens = LensGalaxyParameterDistribution()
-        >>> print(lens.lens_functions)
-        >>> print(lens.lens_priors)
-        >>> print(lens.lens_priors_params)
-    directory : `str`
-        directory to store the interpolators
-        default: './interpolator_pickle'
-    **kwargs : 
-        keyword arguments to pass to the parent classes
-
-    Examples
-    --------
-    >>> from ler.lens_galaxy_population import LensGalaxyParameterDistribution
-    >>> lens = LensGalaxyParameterDistribution()
-    >>> lensed_params = lens.sample_lens_parameters(size=1000)
-    >>> lensed_params.keys()
-
-    Instance Attributes
-    ----------
-    LensGalaxyPopulation class has the following instance attributes:\n
-    +-------------------------------------+----------------------------------+
-    | Atrributes                          | Type                             |
-    +=====================================+==================================+
-    |:attr:`~npool`                       | `int`                            |
-    +-------------------------------------+----------------------------------+
-    |:attr:`~z_min`                       | `float`                          |
-    +-------------------------------------+----------------------------------+
-    |:attr:`~z_max`                       | `float`                          |
-    +-------------------------------------+----------------------------------+
-    |:attr:`~cosmo`                       | `astropy.cosmology`              |
-    +-------------------------------------+----------------------------------+
-    |:attr:`~event_type`                  | `str`                            |
-    +-------------------------------------+----------------------------------+
-    |:attr:`~directory`                   | `str`                            |
-    +-------------------------------------+----------------------------------+
-    |:attr:`~create_new_interpolator`     | `dict`                           |
-    +-------------------------------------+----------------------------------+
-    |:attr:`~lens_param_samplers`         | `dict`                           |
-    +-------------------------------------+----------------------------------+
-    |:attr:`~lens_param_samplers_params`  | `dict`                           |
-    +-------------------------------------+----------------------------------+
-    |:attr:`~lens_sampler_names`          | `dict`                           |
-    +-------------------------------------+----------------------------------+
-    |:attr:`~lens_functions`              | `dict`                           |
-    +-------------------------------------+----------------------------------+
-    |:attr:`~normalization_pdf_z_lensed`  | `float`                          |
-    +-------------------------------------+----------------------------------+
-
-    Instance Methods
-    ----------
-    LensGalaxyPopulation class has the following instance methods:\n
-    +-------------------------------------+----------------------------------+
-    | Methods                             | Type                             |
-    +=====================================+==================================+
-    |:meth:`~sample_lens_parameters`      | Function to call the specific    |
-    |                                     | galaxy lens parameters sampler   |
-    |                                     | routine.                         |
-    +-------------------------------------+----------------------------------+
-    |:meth:`~sample_all_routine`          | Function to sample galaxy lens   |
-    |                                     | parameters along with the source |
-    |                                     | parameters.                      |
-    +-------------------------------------+----------------------------------+
-    |:meth:`~strongly_lensed_source_redshifts`                               |
-    +-------------------------------------+----------------------------------+
-    |                                     | Function to sample source        |
-    |                                     | redshifts conditioned on the     |
-    |                                     | source being strongly lensed     |
-    +-------------------------------------+----------------------------------+
-    |:meth:`~source_parameters`           | Function to sample gw source     |
-    |                                     | parameters                       |
-    +-------------------------------------+----------------------------------+
-    |:meth:`~lens_redshift_SDSS_catalogue`| Function to sample lens          |
-    |                                     | redshifts, conditioned on the    |
-    |                                     | lens being strongly lensed       |
-    +-------------------------------------+----------------------------------+
-    |:meth:`~axis_rotation_angle_uniform` | Function to sample the axis      |
-    |                                     | rotation angle of the elliptical |
-    |                                     | lens galaxy from a uniform       |
-    |                                     | distribution                     |
-    +-------------------------------------+----------------------------------+
-    |:meth:`~shear_norm`                  | Function to sample the           |
-    |                                     | elliptical lens galaxy shear     |
-    |                                     | from a normal distribution       |
-    +-------------------------------------+----------------------------------+
-    |:meth:`~mass_density_spectral_index_normal`                             |
-    +-------------------------------------+----------------------------------+
-    |                                     | Function to sample the lens      |
-    |                                     | galaxy spectral index of the     |
-    |                                     | mass density profile from a      |
-    |                                     | normal distribution              |
-    +-------------------------------------+----------------------------------+
-    |:meth:`~compute_einstein_radii`      | Function to compute the Einstein |
-    |                                     | radii of the lens galaxies       |
-    +-------------------------------------+----------------------------------+
-    |:meth:`~rjs_with_cross_section_SIE`  | Function to conduct rejection    |
-    |                                     | sampling wrt einstein radius     |
-    +-------------------------------------+----------------------------------+
-    |:meth:`~rjs_with_cross_section_SIE`  | Function to conduct rejection    |
-    |                                     | sampling wrt cross_section       |
-    +-------------------------------------+----------------------------------+
-    |:attr:`~rejection_sample_sl`         | Function to conduct rejection    |
-    |                                     | sampling with the given rejection|
-    |                                     | sampling function                |
-    +-------------------------------------+----------------------------------+
-    |:attr:`~sample_source_redshift_sl`   | Function to sample source        |
-    |                                     | redshifts conditioned on the     |
-    |                                     | source being strongly lensed     |
-    +-------------------------------------+----------------------------------+
-    |:attr:`~sample_lens_redshift`        | Function to sample lens          |
-    |                                     | redshifts, conditioned on the    |
-    |                                     | lens being strongly lensed       |
-    +-------------------------------------+----------------------------------+
-    |:attr:`~sample_axis_rotation_angle`  | Function to sample the axis      |
-    |                                     | rotation angle of the elliptical |
-    |                                     | lens galaxy from a uniform       |
-    |                                     | distribution                     |
-    +-------------------------------------+----------------------------------+
-    |:attr:`~sample_shear`                | Function to sample the           |
-    |                                     | elliptical lens galaxy shear     |
-    |                                     | from a normal distribution       |
-    +-------------------------------------+----------------------------------+
-    |:attr:`~sample_mass_density_spectral_index`                             |
-    +-------------------------------------+----------------------------------+
-    |                                     | Function to sample the lens      |
-    |                                     | galaxy spectral index of the     |
-    |                                     | mass density profile from a      |
-    |                                     | normal distribution              |
-    +-------------------------------------+----------------------------------+
+        >>> lensed_params = lens.sample_lens_parameters(size=1000)
+        >>> lensed_params.keys()
+
+        Instance Attributes
+        ----------
+        LensGalaxyPopulation class has the following instance attributes:\n
+        +-------------------------------------+----------------------------------+
+        | Atrributes                          | Type                             |
+        +=====================================+==================================+
+        |:attr:`~npool`                       | `int`                            |
+        +-------------------------------------+----------------------------------+
+        |:attr:`~z_min`                       | `float`                          |
+        +-------------------------------------+----------------------------------+
+        |:attr:`~z_max`                       | `float`                          |
+        +-------------------------------------+----------------------------------+
+        |:attr:`~cosmo`                       | `astropy.cosmology`              |
+        +-------------------------------------+----------------------------------+
+        |:attr:`~event_type`                  | `str`                            |
+        +-------------------------------------+----------------------------------+
+        |:attr:`~directory`                   | `str`                            |
+        +-------------------------------------+----------------------------------+
+        |:attr:`~create_new_interpolator`     | `dict`                           |
+        +-------------------------------------+----------------------------------+
+        |:attr:`~lens_param_samplers`         | `dict`                           |
+        +-------------------------------------+----------------------------------+
+        |:attr:`~lens_param_samplers_params`  | `dict`                           |
+        +-------------------------------------+----------------------------------+
+        |:attr:`~lens_sampler_names`          | `dict`                           |
+        +-------------------------------------+----------------------------------+
+        |:attr:`~lens_functions`              | `dict`                           |
+        +-------------------------------------+----------------------------------+
+        |:attr:`~normalization_pdf_z_lensed`  | `float`                          |
+        +-------------------------------------+----------------------------------+
+
+        Instance Methods
+        ----------
+        LensGalaxyPopulation class has the following instance methods:\n
+        +-------------------------------------+----------------------------------+
+        | Methods                             | Type                             |
+        +=====================================+==================================+
+        |:meth:`~sample_lens_parameters`      | Function to call the specific    |
+        |                                     | galaxy lens parameters sampler   |
+        |                                     | routine.                         |
+        +-------------------------------------+----------------------------------+
+        |:meth:`~sample_all_routine_sie`      | Function to sample galaxy lens   |
+        |                                     | parameters along with the source |
+        |                                     | parameters.                      |
+        +-------------------------------------+----------------------------------+
+        |:meth:`~strongly_lensed_source_redshifts`                               |
+        +-------------------------------------+----------------------------------+
+        |                                     | Function to sample source        |
+        |                                     | redshifts conditioned on the     |
+        |                                     | source being strongly lensed     |
+        +-------------------------------------+----------------------------------+
+        |:meth:`~source_parameters`           | Function to sample gw source     |
+        |                                     | parameters                       |
+        +-------------------------------------+----------------------------------+
+        |:meth:`~lens_redshift_SDSS_catalogue`| Function to sample lens          |
+        |                                     | redshifts, conditioned on the    |
+        |                                     | lens being strongly lensed       |
+        +-------------------------------------+----------------------------------+
+        |:meth:`~axis_rotation_angle_uniform` | Function to sample the axis      |
+        |                                     | rotation angle of the elliptical |
+        |                                     | lens galaxy from a uniform       |
+        |                                     | distribution                     |
+        +-------------------------------------+----------------------------------+
+        |:meth:`~shear_norm`                  | Function to sample the           |
+        |                                     | elliptical lens galaxy shear     |
+        |                                     | from a normal distribution       |
+        +-------------------------------------+----------------------------------+
+        |:meth:`~density_profile_slope_normal`                             |
+        +-------------------------------------+----------------------------------+
+        |                                     | Function to sample the lens      |
+        |                                     | galaxy spectral index of the     |
+        |                                     | mass density profile from a      |
+        |                                     | normal distribution              |
+        +-------------------------------------+----------------------------------+
+        |:meth:`~compute_einstein_radii`      | Function to compute the Einstein |
+        |                                     | radii of the lens galaxies       |
+        +-------------------------------------+----------------------------------+
+        |:meth:`~rjs_with_cross_section_sis`  | Function to conduct rejection    |
+        |                                     | sampling wrt einstein radius     |
+        +-------------------------------------+----------------------------------+
+        |:meth:`~rjs_with_cross_section_sie`  | Function to conduct rejection    |
+        |                                     | sampling wrt cross_section       |
+        +-------------------------------------+----------------------------------+
+        |:attr:`~rejection_sample_sl`         | Function to conduct rejection    |
+        |                                     | sampling with the given rejection|
+        |                                     | sampling function                |
+        +-------------------------------------+----------------------------------+
+        |:attr:`~sample_source_redshift_sl`   | Function to sample source        |
+        |                                     | redshifts conditioned on the     |
+        |                                     | source being strongly lensed     |
+        +-------------------------------------+----------------------------------+
+        |:attr:`~sample_lens_redshift`        | Function to sample lens          |
+        |                                     | redshifts, conditioned on the    |
+        |                                     | lens being strongly lensed       |
+        +-------------------------------------+----------------------------------+
+        |:attr:`~sample_axis_rotation_angle`  | Function to sample the axis      |
+        |                                     | rotation angle of the elliptical |
+        |                                     | lens galaxy from a uniform       |
+        |                                     | distribution                     |
+        +-------------------------------------+----------------------------------+
+        |:attr:`~sample_shear`                | Function to sample the           |
+        |                                     | elliptical lens galaxy shear     |
+        |                                     | from a normal distribution       |
+        +-------------------------------------+----------------------------------+
+        |:attr:`~sample_density_profile_slope`                             |
+        +-------------------------------------+----------------------------------+
+        |                                     | Function to sample the lens      |
+        |                                     | galaxy spectral index of the     |
+        |                                     | mass density profile from a      |
+        |                                     | normal distribution              |
+        +-------------------------------------+----------------------------------+
     """
 
     # Attributes
     cbc_pop = None
-    """:class:`~CompactBinaryPopulation` class\n
-    This is an already initialized class that contains a function (CompactBinaryPopulation.sample_gw_parameters) that actually samples the source parameters. 
+    """:class:`~CBCSourceParameterDistribution` class\n
+    This is an already initialized class that contains a function (CBCSourceParameterDistribution.sample_gw_parameters) that actually samples the source parameters. 
     """
 
     z_min = None
@@ -216,85 +214,51 @@ class LensGalaxyParameterDistribution(CBCSourceParameterDistribution, ImagePrope
         z_max=10.0,
         cosmology=None,
         event_type="BBH",
-        lens_type="epl_galaxy",
+        lens_type="epl_shear_galaxy",
         lens_functions= None,
-        lens_priors=None,
-        lens_priors_params=None,
+        lens_functions_params=None,
+        lens_param_samplers=None,
+        lens_param_samplers_params=None,
         directory="./interpolator_pickle",
         create_new_interpolator=False,
-        **kwargs
+        buffer_size=1000,
+        **kwargs  # for initialization of CBCSourceParameterDistribution class and ImageProperties class
     ):
-        
-        self.npool = npool
-        self.z_min = z_min
-        self.z_max = z_max
-        self.cosmo = cosmology if cosmology else LambdaCDM(H0=70, Om0=0.3, Ode0=0.7)
-        self.event_type = event_type
-        self.directory = directory
-        # initialize the interpolator's parameters
-        self.create_new_interpolator = dict(
-            redshift_distribution=dict(create_new=False, resolution=1000),
-            z_to_luminosity_distance=dict(create_new=False, resolution=1000),
-            velocity_dispersion=dict(create_new=False, resolution=1000),
-            axis_ratio=dict(create_new=False, resolution=1000),
-            optical_depth=dict(create_new=False, resolution=200),
-            z_to_Dc=dict(create_new=False, resolution=1000),
-            Dc_to_z=dict(create_new=False, resolution=1000),
-            angular_diameter_distance=dict(create_new=False, resolution=1000),
-            differential_comoving_volume=dict(create_new=False, resolution=1000),
-            luminosity_distance_to_z=dict(create_new=False, resolution=1000),
-        )
-        if isinstance(create_new_interpolator, dict):
-            self.create_new_interpolator.update(create_new_interpolator)
-        elif create_new_interpolator is True:
-            self.create_new_interpolator = dict(
-                redshift_distribution=dict(create_new=True, resolution=1000),
-                z_to_luminosity_distance=dict(create_new=True, resolution=1000),
-                velocity_dispersion=dict(create_new=True, resolution=1000),
-                axis_ratio=dict(create_new=True, resolution=1000),
-                optical_depth=dict(create_new=True, resolution=200),
-                z_to_Dc=dict(create_new=True, resolution=1000),
-                Dc_to_z=dict(create_new=True, resolution=1000),
-                angular_diameter_distance=dict(create_new=True, resolution=1000),
-                differential_comoving_volume=dict(create_new=True, resolution=1000),
-                luminosity_distance_to_z=dict(create_new=True, resolution=1000),
-            )
-
-        # dealing with prior functions and categorization
-        self.lens_param_samplers, self.lens_param_samplers_params, self.lens_sampler_names, self.lens_functions = self.lens_priors_categorization(lens_type, lens_priors,
-        lens_priors_params, lens_functions)
-        
-        # function initialization
-        self.sample_lens_parameters_routine = getattr(self, self.lens_functions['param_sampler_type'])
-        self.rejection_sample_sl = getattr(self, self.lens_functions['strong_lensing_condition'])
+        print("\nInitializing LensGalaxyParameterDistribution class...\n")          
+        self.event_type = event_type  # needed for the source population
+        self.buffer_size = buffer_size  # buffer size for sampling lens parameters        
 
         # initializing parent classes
-        self.class_initialization_lens(params=kwargs);
-
-        # initializing samplers
-        # self.sample_velocity_dispersion and self.sample_axis_ratio are initialized in OpticalDepth class
-        self.sample_source_redshift_sl = self.lens_param_samplers["source_redshift_sl"]
-        self.sample_lens_redshift = self.lens_param_samplers["lens_redshift"]
-        # self.sample_axis_ratio = self.lens_param_samplers["axis_ratio"]
-        self.sample_axis_rotation_angle = self.lens_param_samplers[
-            "axis_rotation_angle"
-        ]
-        self.sample_shear = self.lens_param_samplers["shear"]
-        self.sample_mass_density_spectral_index = self.lens_param_samplers[
-            "mass_density_spectral_index"
-        ]
-        self.sample_source_parameters = self.lens_param_samplers["source_parameters"]
-
-        # extra care to the velocity dispersion sampler
-        if self.lens_param_samplers["velocity_dispersion"]  == "velocity_dispersion_ewoud":
-            vd_inv_cdf = self.vd_inv_cdf
-            zl_list = self.zl_list
-            self.sample_velocity_dispersion = lambda size, zl: velocity_dispersion_z_dependent(size=size, zl=zl, zl_list=zl_list, vd_inv_cdf=vd_inv_cdf)
+        self.class_initialization_lens(
+            npool,
+            z_min,
+            z_max,
+            cosmology,
+            lens_type,
+            lens_functions,
+            lens_functions_params,
+            lens_param_samplers, 
+            lens_param_samplers_params, 
+            directory,
+            create_new_interpolator,
+            params=kwargs,  # related parameters for CBCSourceParameterDistribution, ImageProperties classes
+        );
+
+        # function to sample source redshifts conditioned on the source being strongly lensed 
+        self.sample_source_redshift_sl = getattr(self,self.lens_param_samplers["source_redshift_sl"])
+
+        # interpolated cross section function is not very accurate for rejection sampling
+        if self.lens_functions["cross_section"] == "interpolated_cross_section_function":
+            self.cross_section = self.cross_section_caustic_area
+        # function to sample lens parameters
+        self.sample_lens_parameters_routine = getattr(self, self.lens_functions['param_sampler_type'])  
+        # function to rejection sample wrt lens cross section
+        self.rejection_sample_sl = getattr(self, self.lens_functions['strong_lensing_condition']) 
 
         # To find the normalization constant of the pdf p(z)
         # this under the assumption that the event is strongly lensed
         # Define the merger-rate density function
-        pdf_unnormalized_ = lambda z: self.merger_rate_density_detector_frame(np.array([z]), param=self.merger_rate_density_param) * self.strong_lensing_optical_depth(np.array([z]))
+        pdf_unnormalized_ = lambda z: self.merger_rate_density_detector_frame(np.array([z])) * self.optical_depth.function(np.array([z]))
         pdf_unnormalized = lambda z: pdf_unnormalized_(z)[0]
 
         self.normalization_pdf_z_lensed = quad(
@@ -303,17 +267,58 @@ class LensGalaxyParameterDistribution(CBCSourceParameterDistribution, ImagePrope
             self.z_max
         )[0]
 
-    def class_initialization_lens(self, params=None):
+    def class_initialization_lens(self, npool, z_min, z_max, cosmology, lens_type, lens_functions, lens_functions_params, lens_param_samplers,  lens_param_samplers_params,  directory, create_new_interpolator, params):
         """
-        Function to initialize the parent classes
+            Initialize the LensGalaxyParameterDistribution class.
+
+            Parameters
+            ----------
+            npool : `int`
+                number of processors to use for sampling
+            z_min : `float`
+                minimum redshift of the lens galaxy
+            z_max : `float`
+                maximum redshift of the lens galaxy
+            cosmology : `astropy.cosmology`
+                cosmology object
+            lens_type : `str`
+                type of the lens galaxy
+            lens_functions : `dict`
+                dictionary with the lens related functions
+            lens_functions_params : `dict`
+                dictionary with the parameters for the lens related functions
+            lens_param_samplers : `dict`
+                dictionary with the priors for the sampler
+            lens_param_samplers_params : `dict`
+                dictionary with the parameters for the priors of the sampler
+            directory : `str`
+                directory where the interpolators are saved
+            create_new_interpolator : `bool`
+                if True, creates a new interpolator
+            params : `dict`
+                additional parameters for the CBCSourceParameterDistribution and ImageProperties classes
 
-        Parameters
-        ----------
-        params : `dict`
-            dictionary of parameters to initialize the parent classes
         """
 
-        # initialization of CompactBinaryPopulation class
+        # initialize the optical depth class
+        # this also initializes the lens related parameter samplers and functions:
+        # self.lens_param_samplers, self.lens_param_samplers_params, self.lens_functions, self.lens_functions_params
+        OpticalDepth.__init__(
+            self,
+            npool = npool,
+            z_min = z_min,
+            z_max = z_max,
+            cosmology=cosmology,
+            lens_type = lens_type,
+            lens_functions = lens_functions,
+            lens_functions_params = lens_functions_params,
+            lens_param_samplers=lens_param_samplers,
+            lens_param_samplers_params=lens_param_samplers_params,
+            directory=directory,
+            create_new_interpolator=create_new_interpolator,
+        )
+
+        # initialization of CBCSourceParameterDistribution class
         # it also initializes the CBCSourceRedshiftDistribution class
         # list of relevant initialized instances,
         # 1. self.sample_source_redshift
@@ -340,42 +345,18 @@ class LensGalaxyParameterDistribution(CBCSourceParameterDistribution, ImagePrope
             create_new_interpolator=self.create_new_interpolator,
         )
 
-        # initialize the optical depth class
-        # follwing attributes are initialized
-        # 1. self.strong_lensing_optical_depth
-        # 2. self.sample_velocity_dispersion
-        # 3. self.sample_axis_ratio
-        OpticalDepth.__init__(
-            self,
-            npool=self.npool,
-            z_min=self.z_min,
-            z_max=self.z_max,
-            optical_depth_function=self.lens_functions["optical_depth"],
-            sampler_priors=dict(
-                velocity_dispersion=self.lens_param_samplers["velocity_dispersion"],
-                axis_ratio=self.lens_param_samplers["axis_ratio"],
-            ),
-            sampler_priors_params=dict(
-                velocity_dispersion=self.lens_param_samplers_params[
-                    "velocity_dispersion"],
-                axis_ratio=self.lens_param_samplers_params["axis_ratio"],
-            ),
-            cosmology=self.cosmo,
-            directory=self.directory,
-            create_new_interpolator=self.create_new_interpolator,
-        )
-
         # initialize the image properties class
         input_params_image = dict(
             n_min_images=2,
             n_max_images=4,
-            geocent_time_min=1126259462.4,
-            geocent_time_max=1126259462.4+365*24*3600*10,
+            time_window=365*24*3600*20,
+            # geocent_time_min=1126259462.4,
+            # geocent_time_max=1126259462.4+365*24*3600*2,
             lens_model_list=["EPL_NUMBA", "SHEAR"],
         )
         input_params_image.update(params)
 
-        print("input_params_image", input_params_image)
+        # print("input_params_image", input_params_image)
         ImageProperties.__init__(
             self,
             npool=self.npool,
@@ -385,111 +366,147 @@ class LensGalaxyParameterDistribution(CBCSourceParameterDistribution, ImagePrope
             n_max_images=input_params_image["n_max_images"],
             lens_model_list=input_params_image["lens_model_list"],
             cosmology=self.cosmo,
-            geocent_time_min=input_params_image["geocent_time_min"],
-            geocent_time_max=input_params_image["geocent_time_max"],
+            time_window=input_params_image["time_window"],
+            # geocent_time_min=input_params_image["geocent_time_min"],
+            # geocent_time_max=input_params_image["geocent_time_max"],
             spin_zero=input_params["spin_zero"],
             spin_precession=input_params["spin_precession"],
             directory=self.directory,
             create_new_interpolator=self.create_new_interpolator,
         )
 
-    def lens_priors_categorization(
-        self, lens_type, lens_priors=None, lens_priors_params=None, lens_functions=None,
-    ):
+    def sample_lens_parameters(self, size=1000,):
         """
-        Function to categorize the lens priors/samplers
-
-        Parameters
-        ----------
-            lens_type : `str`
-                lens type
-                e.g. 'epl_galaxy' for elliptical power-law galaxy
-            lens_priors : `dict`
-                dictionary of priors
-            lens_priors_params : `dict`
-                dictionary of priors parameters
-            lens_functions : `dict`
-                dictionary of lens functions
-
-        Returns
-        -------
-        lens_priors_ : `dict`
-            dictionary of priors
-        lens_priors_params_ : `dict`
-            dictionary of priors parameters
-        lens_sampler_names_ : `dict`
-            dictionary of sampler names
-        lens_functions_ : `dict`
-            dictionary of lens functions
+            Function to sample galaxy lens parameters along with the source parameters, conditioned on the source being strongly lensed.
+
+            Parameters
+            ----------
+            size : `int`
+                number of lens parameters to sample
+
+            Returns
+            -------
+            lens_parameters : `dict`
+                dictionary of sampled lens parameters and source parameters. \n
+                keys: \n
+                zl: lens redshifts \n
+                zs: source redshifts, lensed condition applied\n
+                sigma: velocity dispersions \n
+                q: axis ratios \n
+                theta_E: Einstein radii \n
+                phi: axis rotation angle \n
+                e1: ellipticity component 1 \n
+                e2: ellipticity component 2 \n
+                gamma1: shear component 1 \n
+                gamma2: shear component 2 \n
+                gamma: density profile slope distribution \n
+                geocent_time: time of arrival of the unlensed signal\n
+                phase: phase of the unlensed signal\n
+                psi: polarization angle of the unlensed signal\n
+                theta_jn: inclination angle of the unlensed signal\n
+                luminosity_distance: luminosity distance of the source\n
+                mass_1_source: mass 1 (larger) of the source\n
+                mass_2_source: mass 2 (smaller) of the source\n
+                ra: right ascension of the source\n
+                dec: declination of the source\n
+
+            Examples
+            --------
+            >>> from ler.lens_galaxy_population import LensGalaxyParameterDistribution
+            >>> od = LensGalaxyParameterDistribution(lens_param_samplers=dict(velocity_dispersion="velocity_dispersion_ewoud"))
+            >>> print(od.sample_lens_parameters(size=10))
         """
 
-        if lens_type == "epl_galaxy":
-            lens_priors_ = dict(
-                source_redshift_sl="strongly_lensed_source_redshifts",
-                lens_redshift="lens_redshift_SDSS_catalogue",
-                velocity_dispersion="velocity_dispersion_ewoud",
-                axis_ratio="axis_ratio_rayleigh",
-                axis_rotation_angle="axis_rotation_angle_uniform",
-                shear="shear_norm",
-                mass_density_spectral_index="mass_density_spectral_index_normal",
-                source_parameters="sample_gw_parameters",
-            )
-            lens_priors_params_ = dict(
-                source_redshift_sl=None,
-                lens_redshift=None,
-                velocity_dispersion=None,
-                axis_ratio=dict(q_min=0.2, q_max=1.),
-                axis_rotation_angle=dict(phi_min=0.0, phi_max=2 * np.pi),
-                shear=dict(scale=0.05),
-                mass_density_spectral_index=dict(mean=2.0, std=0.2),
-                source_parameters=None,
-            )
-            lens_functions_ = dict(
-                strong_lensing_condition="rjs_with_cross_section_SIE",
-                optical_depth="optical_depth_SIE_hemanta",
-                param_sampler_type="sample_all_routine",
-            )
-        else:
-            raise ValueError("lens_type not recognized")
-
-        # update the priors if input is given
-        if lens_priors:
-            lens_priors_.update(lens_priors)
-        if lens_priors_params:
-            lens_priors_params_.update(lens_priors_params)
-        if lens_functions:
-            lens_functions_.update(lens_functions)
-
-        # dict of sampler names with description
-        lens_sampler_names_ = dict(
-            sample_source_redshift_sl="source parameters conditioned on the source being strongly lensed",
-            sample_lens_redshift="lens redshift",
-            sample_velocity_dispersion="velocity dispersion of elliptical galaxy",
-            sample_axis_ratio="axis ratio of elliptical galaxy",
-            sample_axis_rotation_angle="axis rotation angle of elliptical galaxy    ",
-            sample_shear="shear of elliptical galaxy",
-            sample_mass_density_spectral_index="mass density spectral index of elliptical power-law galaxy",
-            sample_source_parameters="source parameters other than redshift",
+        print(f"sampling lens parameters with {self.lens_functions['param_sampler_type']}...")
+
+        # get sample lens parameters, strongly lensed
+        lens_parameters = self.sample_lens_parameters_routine(
+            size=size
         )
 
-        return(lens_priors_, lens_priors_params_, lens_sampler_names_, lens_functions_)
+        # sample gravitional waves source parameter
+        param = dict(zs=lens_parameters["zs"])
+        # zs won't be sampling again
+        gw_param = self.sample_gw_parameters(size=size, param=param)
+        # Add source params strongly lensed to the lens params
+        lens_parameters.update(gw_param)
 
-    def sample_lens_parameters(
-        self,
-        size=1000,
-        lens_parameters_input=None,
-    ):
+        return lens_parameters
+    
+    def sample_all_routine_sis_sl(self, size=1000):
         """
-        Function to call the specific galaxy lens parameters sampler routine.
+        Function to sample galaxy lens parameters. SIS cross section is used for rejection sampling.
         """
 
-        return self.sample_lens_parameters_routine(
-            size=size, lens_parameters_input=lens_parameters_input
-        )
+        buffer_size = self.buffer_size
+
+        # Sample source redshifts from the source population
+        # rejection sampled with optical depth
+        zs = self.sample_source_redshift_sl(size)
 
-    def sample_all_routine(self, size=1000, lens_parameters_input=None):
+        # # Sample lens redshifts
+        zl = self.lens_redshift.rvs(size, zs)
+
+        sigma = np.zeros(size)
+        theta_E = np.zeros(size)
+        
+        sigma_max = self.velocity_dispersion.info['sigma_max']
+
+        for i in tqdm(range(size), ncols=100, disable=False):
+            zs_ = zs[i]*np.ones(buffer_size)
+            zl_ = zl[i]*np.ones(buffer_size)
+
+            # cross_section_max calculation
+            theta_E_max = self.compute_einstein_radii(np.array([sigma_max]), np.array([zl[i]]), np.array([zs[i]]))[0]
+            cross_section_max = np.pi*theta_E_max**2
+
+            while True:
+                # Create a dictionary of the lens parameters; sigma, theta_E, q, phi, e1, e2
+                lens_parameters_ = self.sampling_routine_sie_nsl(zl_, zs_, size=buffer_size)
+
+                # Rejection sample based on the lensing probability, that is, rejection sample wrt theta_E
+                lens_parameters_, mask, cross_section_max_ = self.rjs_with_cross_section_sis(
+                    lens_parameters_, cross_section_max
+                )  # proportional to pi theta_E^2
+
+                if cross_section_max_>cross_section_max:
+                    cross_section_max = cross_section_max_
+
+                if np.sum(mask) > 0:
+                    break
+            
+            sigma[i] = lens_parameters_["sigma"][0]
+            theta_E[i] = lens_parameters_["theta_E"][0]
+            
+        # sample additional lens parameters
+        # P(q|SL), P(gamma|SL), P(gamma1, gamma2|SL)
+        q = self.axis_ratio.rvs(size)
+        gamma = self.density_profile_slope_sl.rvs(size)
+        gamma1, gamma2 = self.external_shear_sl.rvs(size)
+
+        phi = self.axis_rotation_angle.rvs(size)
+        e1, e2 = phi_q2_ellipticity_hemanta(phi, q)
+
+        # Create a dictionary of the lens parameters
+        lens_parameters = {
+            "zl": zl,
+            "zs": zs,
+            "sigma": sigma,
+            "theta_E": theta_E,
+            "q": q,
+            "phi": phi,
+            "e1": e1,
+            "e2": e2,
+            "gamma": gamma,
+            "gamma1": gamma1,
+            "gamma2": gamma2,
+        }
+
+        return lens_parameters
+
+    def sample_all_routine_sie_sl(self, size=1000):
         """
-        Function to sample galaxy lens parameters along with the source parameters.
+        Function to sample galaxy lens parameters. SIE cross section is used for rejection sampling.
 
         Parameters
         ----------
@@ -512,406 +529,416 @@ class LensGalaxyParameterDistribution(CBCSourceParameterDistribution, ImagePrope
             e2: ellipticity component 2 \n
             gamma1: shear component 1 \n
             gamma2: shear component 2 \n
-            gamma: spectral index of the mass density distribution \n
-            geocent_time: time of arrival of the unlensed signal\n
-            phase: phase of the unlensed signal\n
-            psi: polarization angle of the unlensed signal\n
-            theta_jn: inclination angle of the unlensed signal\n
-            luminosity_distance: luminosity distance of the source\n
-            mass_1_source: mass 1 (larger) of the source\n
-            mass_2_source: mass 2 (smaller) of the source\n
-            ra: right ascension of the source\n
-            dec: declination of the source\n
+            gamma: density profile slope distribution \n
 
         Examples
         --------
         >>> from ler.lens_galaxy_population import LensGalaxyParameterDistribution
         >>> lens = LensGalaxyParameterDistribution()
-        >>> lens.sample_all_routine(size=1000)
+        >>> lens.sample_all_routine_sie(size=1000)
         """
 
-        if lens_parameters_input is None:
-            lens_parameters_input = dict()
-        samplers_params = self.lens_param_samplers_params.copy()
+        buffer_size = self.buffer_size
 
         # Sample source redshifts from the source population
         # rejection sampled with optical depth
-        zs = self.sample_source_redshift_sl(size=size)
-
-        # Sample lens redshifts
-        zl = self.sample_lens_redshift(zs=zs)
-
-        # Sample velocity dispersions
-        try:
-            sigma = self.sample_velocity_dispersion(len(zs))
-        except:
-            sigma = self.sample_velocity_dispersion(len(zs), zl)
-
-        # Sample axis ratios
-        try:
-            q = self.sample_axis_ratio(sigma)
-        except:
-            q = self.sample_axis_ratio(len(sigma))
-
-        # Compute the Einstein radii
-        theta_E = self.compute_einstein_radii(sigma, zl, zs)
+        zs = self.sample_source_redshift_sl(size)
+
+        # # Sample lens redshifts
+        zl = self.lens_redshift.rvs(size, zs)
+
+        sigma = np.zeros(size)
+        theta_E = np.zeros(size)
+        q = np.zeros(size)
+
+        # # set-up multiprocessing args
+        # q_args_cdf_values = self.axis_ratio.cdf_values
+        # q_args_x_array = self.axis_ratio.x_array
+        # q_args_conditioned_y_array = self.axis_ratio.conditioned_y_array
+
+        # sigma_args_cdf_values = self.velocity_dispersion.cdf_values
+        # sigma_args_x_array = self.velocity_dispersion.x_array
+        # sigma_args_conditioned_y_array = self.velocity_dispersion.conditioned_y_array
+
+        # da_args_function_spline = self.angular_diameter_distance.function_spline
+        # da_args_x_array = self.angular_diameter_distance.x_array
+
+        # idx = np.arange(size)
+        # input_params = np.array([(idx[i], zs[i], zl[i], sigma_args_cdf_values, sigma_args_x_array, sigma_args_conditioned_y_array, q_args_cdf_values, q_args_x_array, q_args_conditioned_y_array, da_args_function_spline, da_args_x_array, buffer_size) for i in range(size)], dtype=object)
+
+        # with Pool(processes=self.npool) as pool:
+        #     for result in tqdm(
+        #         pool.imap_unordered(rjs_sie_mp, input_params),
+        #         total=size,
+        #         ncols=100,
+        #         disable=False,
+        #     ):
+        #         (
+        #             idx_,
+        #             sigma_,
+        #             theta_E_,
+        #             q_,
+        #         ) = result
+
+        #         sigma[idx_] = sigma_
+        #         theta_E[idx_] = theta_E_
+        #         q[idx_] = q_
+        # phi = self.axis_rotation_angle.rvs(size)
+        # e1, e2 = phi_q2_ellipticity_hemanta(phi, q)
+        
+        sigma_max = self.velocity_dispersion.info['sigma_max']
+
+        for i in tqdm(range(size), ncols=100, disable=False):
+            zs_ = zs[i]*np.ones(buffer_size)
+            zl_ = zl[i]*np.ones(buffer_size)
+
+            # cross_section_max calculation
+            theta_E_max = self.compute_einstein_radii(np.array([sigma_max]), np.array([zl[i]]), np.array([zs[i]]))[0]
+            sie_factor=1.
+            cross_section_max = sie_factor*np.pi*theta_E_max**2
+
+            while True:
+                # Create a dictionary of the lens parameters; sigma, theta_E, q, phi, e1, e2
+                lens_parameters_ = self.sampling_routine_sie_nsl(zl_, zs_, size=buffer_size)
+
+                # Rejection sample based on the lensing probability, that is, rejection sample wrt theta_E
+                lens_parameters_, mask, cross_section_max_ = self.rjs_with_cross_section_sie_feixu(
+                    lens_parameters_, cross_section_max
+                )  # proportional to pi theta_E^2
+
+                if cross_section_max_>cross_section_max:
+                    cross_section_max = cross_section_max_
+
+                if np.sum(mask) > 0:
+                    break
+            
+            sigma[i] = lens_parameters_["sigma"][0]
+            theta_E[i] = lens_parameters_["theta_E"][0]
+            q[i] = lens_parameters_["q"][0]
+            
+        # sample additional lens parameters
+        # P(gamma|SL), P(gamma1, gamma2|SL)
+        gamma = self.density_profile_slope_sl.rvs(size)
+        gamma1, gamma2 = self.external_shear_sl.rvs(size)
+
+        phi = self.axis_rotation_angle.rvs(size)
+        e1, e2 = phi_q2_ellipticity_hemanta(phi, q)
 
         # Create a dictionary of the lens parameters
         lens_parameters = {
             "zl": zl,
             "zs": zs,
             "sigma": sigma,
-            "q": q,
             "theta_E": theta_E,
+            "q": q,
+            "phi": phi,
+            "e1": e1,
+            "e2": e2,
+            "gamma": gamma,
+            "gamma1": gamma1,
+            "gamma2": gamma2,
         }
 
-        # Rejection sample based on the lensing probability, that is, rejection sample wrt theta_E
-        lens_parameters = self.rejection_sample_sl(
-            lens_parameters
-        )  # proportional to pi theta_E^2
-
-        # Add the lensing parameter dictionaries together
-        lens_parameters = add_dictionaries_together(
-            lens_parameters, lens_parameters_input
-        )
-
-        # check the size of the lens parameters
-        if len(lens_parameters["zl"]) < size:
-            # Run iteratively until we have the right number of lensing parmaeters
-            # print("current sampled size", len(lens_parameters["zl"]))
-            return self.sample_all_routine(
-                size=size, lens_parameters_input=lens_parameters
-            )
-        else:
-            # Trim dicitionary to right size
-            lens_parameters = trim_dictionary(lens_parameters, size)
-
-            # Sample the axis rotation angle
-            lens_parameters["phi"] = self.sample_axis_rotation_angle(size=size)
-
-            # Transform the axis ratio and the angle, to ellipticities e1, e2, using lenstronomy
-            lens_parameters["e1"], lens_parameters["e2"] = phi_q2_ellipticity(
-                lens_parameters["phi"], lens_parameters["q"]
-            )
-
-            # Sample shears
-            lens_parameters["gamma1"], lens_parameters["gamma2"] = self.sample_shear(
-                size=size)
-
-            # Sample the spectral index of the mass density distribution
-            lens_parameters["gamma"] = self.sample_mass_density_spectral_index(
-                size=size)
+        return lens_parameters
+        
+    def sample_all_routine_epl_shear_sl(self, size=1000):
+        """
+            Function to sample galaxy lens parameters along. EPL shear cross section is used for rejection sampling.
+
+            Parameters
+            ----------
+            size : `int`
+                number of lens parameters to sample
+            lens_parameters_input : `dict`
+                dictionary of lens parameters to sample
+
+            Returns
+            -------
+            lens_parameters : `dict`
+                dictionary of lens parameters and source parameters (lens conditions applied): \n
+                zl: lens redshifts \n
+                zs: source redshifts, lensed condition applied\n
+                sigma: velocity dispersions \n
+                q: axis ratios \n
+                theta_E: Einstein radii \n
+                phi: axis rotation angle \n
+                e1: ellipticity component 1 \n
+                e2: ellipticity component 2 \n
+                gamma1: shear component 1 \n
+                gamma2: shear component 2 \n
+                gamma: density profile slope distribution \n
+
+            Examples
+            --------
+            >>> from ler.lens_galaxy_population import LensGalaxyParameterDistribution
+            >>> lens = LensGalaxyParameterDistribution()
+            >>> lens.sample_all_routine_sie(size=1000)
+        """
 
-            # sample gravitional waves source parameter
-            param = dict(zs=lens_parameters["zs"])
-            if samplers_params["source_parameters"]:
-                param.update(self.sample_gw_parameters(size=size))
-            gw_param = self.sample_source_parameters(size=size, param=param)
+        buffer_size = self.buffer_size
 
-            # Add source params strongly lensed to the lens params
-            lens_parameters.update(gw_param)
+        # Sample source redshifts from the source population
+        # rejection sampled with optical depth
+        zs = self.sample_source_redshift_sl(size)
+
+        # # Sample lens redshifts
+        zl = self.lens_redshift.rvs(size, zs)
+
+        sigma = np.zeros(size)
+        theta_E = np.zeros(size)
+        q = np.zeros(size)
+        phi = np.zeros(size)
+        gamma1 = np.zeros(size)
+        gamma2 = np.zeros(size)
+        gamma = np.zeros(size)
+
+        sigma_max = self.velocity_dispersion.info['sigma_max']
+
+        # this will take some time
+        for i in tqdm(range(size), ncols=100, disable=False):
+            zs_ = zs[i]*np.ones(buffer_size)
+            zl_ = zl[i]*np.ones(buffer_size)
+
+            # cross_section_max calculation            
+            theta_E_max = self.compute_einstein_radii(np.array([sigma_max]), np.array([zl[i]]), np.array([zs[i]]))[0]
+            epl_factor=5.
+            cross_section_max = epl_factor*np.pi*theta_E_max**2
+
+            while True:
+                # Create a dictionary of the lens parameters, gamma, gamma1, gamma2, sigma, theta_E, q, phi, e1, e2
+                lens_parameters_ = self.sampling_routine_epl_shear_nsl(zl_, zs_, size=buffer_size)
+
+                # Rejection sample based on the lensing probability, that is, rejection sample wrt theta_E
+                lens_parameters_, mask, cross_section_max_ = self.rejection_sample_sl(
+                    lens_parameters_, cross_section_max
+                )  # proportional to pi theta_E^2
+                
+                if cross_section_max_>cross_section_max:
+                    cross_section_max = cross_section_max_
+
+                # check the size of the lens parameters
+                if np.sum(mask) > 0:
+                    break
+
+            sigma[i] = lens_parameters_["sigma"][0]
+            theta_E[i] = lens_parameters_["theta_E"][0]
+            q[i] = lens_parameters_["q"][0]
+            phi[i] = lens_parameters_["phi"][0]
+            gamma[i] = lens_parameters_["gamma"][0]
+            gamma1[i] = lens_parameters_["gamma1"][0]
+            gamma2[i] = lens_parameters_["gamma2"][0]
+        e1, e2 = phi_q2_ellipticity_hemanta(phi, q)
 
-            return lens_parameters
+        # Create a dictionary of the lens parameters
+        lens_parameters = {
+            "zl": zl,
+            "zs": zs,
+            "sigma": sigma,
+            "theta_E": theta_E,
+            "q": q,
+            "phi": phi,
+            "e1": e1,
+            "e2": e2,
+            "gamma": gamma,
+            "gamma1": gamma1,
+            "gamma2": gamma2,
+        }
 
-    def strongly_lensed_source_redshifts(self, size=1000):
+        return lens_parameters
+        
+    def sampling_routine_sis_nsl(self, zl, zs, size=1000):
         """
-        Function to sample source redshifts and other parameters, conditioned on the source being strongly lensed.
+        Function to sample SIS lens related parameters.
 
         Parameters
         ----------
+        zl : `float`
+            lens redshifts
+        zs : `float`
+            source redshifts
         size : `int`
             number of lens parameters to sample
 
         Returns
         -------
-        redshifts : `float`
-            source redshifts conditioned on the source being strongly lensed
-
-        Examples
-        --------
-        >>> from ler.lens_galaxy_population import LensGalaxyParameterDistribution
-        >>> lens = LensGalaxyParameterDistribution()
-        >>> lens.strongly_lensed_source_redshifts(size=1000)
-        """
-
-        z_max = self.z_max
-
-        def zs_function(zs_sl):
-            # get zs
-            # self.sample_source_redshifts from CBCSourceRedshiftDistribution class
-            zs = self.sample_zs(size)  # this function is from CompactBinaryPopulation class
-            # put strong lensing condition with optical depth
-            tau = self.strong_lensing_optical_depth(zs)
-            tau_max = self.strong_lensing_optical_depth(np.array([z_max]))[0] # tau increases with z
-            r = np.random.uniform(0, tau_max, size=len(zs)) 
-            # Add the strongly lensed source redshifts to the list
-            # pick strongly lensed sources
-            zs_sl += list(zs[r < tau])  # list concatenation
-
-            # Check if the zs_sl are larger than requested size
-            if len(zs_sl) >= size:
-                # Trim list to right size
-                zs_sl = zs_sl[:size]
-                return zs_sl
-            else:
-                # Run iteratively until we have the right number of lensing parmaeters
-                return zs_function(zs_sl)
-
-        zs_sl = []
-
-        return np.array(zs_function(zs_sl))
-
-    def source_parameters(self, size, get_attribute=False, param=None):
+        lens_parameters : `dict`
+            dictionary of sampled lens parameters.
+            keys: sigma, theta_E
         """
-        Function to sample gw source parameters
-
-        Parameters
-        ----------
-        size : `int`
-            Number of samples to draw
-        param : `dict`
-            Allows to pass in parameters as dict.
-            param =
+        try:
+            sigma = self.velocity_dispersion.rvs(size, zl)
+        except:
+            sigma = self.velocity_dispersion.rvs(size)
 
-        Returns
-        ----------
-        source_parameters : `dict`
-            Dictionary of source parameters
-            source_parameters.keys() = ['mass_1', 'mass_2', 'mass_1_source', 'mass_2_source', 'zs', 'luminosity_distance', 'inclination', 'polarization_angle', 'phase', 'geocent_time', 'ra', 'dec', 'a_1', 'a_2', 'tilt_1', 'tilt_2', 'phi_12', 'phi_jl']
+        # Compute the Einstein radii
+        theta_E = self.compute_einstein_radii(sigma, zl, zs)
 
-        Examples
-        --------
-        >>> from ler.lens_galaxy_population import LensGalaxyParameterDistribution
-        >>> lens = LensGalaxyParameterDistribution()
-        >>> lens.source_parameters(size=1000)
-        """
+        # Create a dictionary of the lens parameters
+        lens_parameters = {
+            "zs": zs,
+            "zl": zl,
+            "sigma": sigma,
+            "theta_E": theta_E,
+        }
 
-        if get_attribute:
-            return lambda size: self.sample_gw_parameters(size=size, param=param)
-        else:
-            # sample gravitional waves source parameter
-            return self.sample_gw_parameters(size=size, param=param)
+        return lens_parameters
     
-    def lens_redshift_SDSS_catalogue(self, zs, get_attribute=False, param=None):
+    def sampling_routine_sie_nsl(self, zl, zs, size=1000):
         """
-        Function to sample lens redshifts, conditioned on the lens being strongly lensed
+        Function to sample SIE lens related parameters.
 
         Parameters
         ----------
+        zl : `float`
+            lens redshifts
         zs : `float`
             source redshifts
-        get_attribute : `bool`
-            If True, returns a function that can be called with zs as input
+        size : `int`
+            number of lens parameters to sample
 
         Returns
         -------
-        zl : `float`
-            lens redshifts
-
-        Examples
-        --------
-        >>> from ler.lens_galaxy_population import LensGalaxyParameterDistribution
-        >>> lens = LensGalaxyParameterDistribution()
-        >>> lens.lens_redshift_SDSS_catalogue(zs=1.0)
-        """
-
-        splineDc = self.splineDc  # spline coefficients for the comoving distance and redshifts
-        splineDcInv = self.splineDcInv  # spline coefficients for the redshifts and comoving distance
-        u = np.linspace(0, 1, 500)
-        cdf = (10 * u**3 - 15 * u**4 + 6 * u**5)  # See the integral of Eq. A7 of https://arxiv.org/pdf/1807.07062.pdf (cdf)
-        zs = np.array([zs]).reshape(-1)
-
-        # lens redshifts
-        #return self.Dc_to_z(lens_galaxy_Dc)
-        if get_attribute:
-            return njit(lambda zs: lens_redshift_SDSS_catalogue(zs, splineDc, splineDcInv, u, cdf))
-        else:
-            return lens_redshift_SDSS_catalogue(zs, splineDc, splineDcInv, u, cdf)
-
-    def axis_rotation_angle_uniform(
-        self, size=1000, phi_min=0.0, phi_max=2 * np.pi, get_attribute=False,param=None
-    ):
+        lens_parameters : `dict`
+            dictionary of sampled lens parameters.
+            keys: sigma, q, phi
         """
-        Function to sample the axis rotation angle of the elliptical lens galaxy from a uniform distribution.
 
-        Parameters
-        ----------
-        size : `int`
-            number of lens parameters to sample
-        phi_min : `float`
-            minimum axis rotation angle of the elliptical lens galaxy
-        phi_max : `float`
-            maximum axis rotation angle of the elliptical lens galaxy
-        get_attribute : `bool`
-            If True, returns a function that can be called with size as input
-        param : `dict`
-            Allows to pass in above parameters as dict.
-            e.g. param = dict(phi_min=0.0, phi_max=2 * np.pi)
+        lens_parameters = self.sampling_routine_sis_nsl(zl, zs, size=size)
 
-        Returns
-        -------
-        phi : `float`
-            axis rotation angle of the elliptical lens galaxy
+        # Sample axis ratios
+        try:
+            q = self.axis_ratio.rvs(size, lens_parameters["sigma"])
+        except:
+            q = self.axis_ratio.rvs(size)
 
-        Examples
-        --------
-        >>> from ler.lens_galaxy_population import LensGalaxyParameterDistribution
-        >>> lens = LensGalaxyParameterDistribution()
-        >>> lens.axis_rotation_angle_uniform(size=1000)
-        """
+        # Sample the axis rotation angle
+        phi = self.axis_rotation_angle.rvs(size)
 
-        if param:
-            phi_min = param["phi_min"]
-            phi_max = param["phi_max"]
+        # Transform the axis ratio and the angle, to ellipticities e1, e2, using lenstronomy
+        # e1, e2 = phi_q2_ellipticity_hemanta(phi, q)
 
-        if get_attribute:
-            return njit(lambda size: np.random.uniform(phi_min, phi_max, size=size))
-        else:
-            # Draw the angles from a uniform distribution
-            return np.random.uniform(phi_min, phi_max, size=size)
+        # Add the lensing parameter dictionaries together
+        lens_parameters["q"] = q
+        lens_parameters["phi"] = phi
+        # lens_parameters["e1"] = e1
+        # lens_parameters["e2"] = e2
 
-    def shear_norm(self, size, scale=0.05, get_attribute=False, param=None):
+        return lens_parameters
+    
+    def sampling_routine_epl_shear_nsl(self, zl, zs, size=1000):
         """
-        Function to sample the elliptical lens galaxy shear from a normal distribution
+        Function to sample EPL and shear related parameters.
 
         Parameters
         ----------
+        zl : `float`
+            lens redshifts
+        zs : `float`
+            source redshifts
         size : `int`
             number of lens parameters to sample
-        scale : `float`
-            standard deviation of the normal distribution
-        get_attribute : `bool`
-            If True, returns a function that can be called with size as input
-        param : `dict`
-            Allows to pass in above parameters as dict.
-            e.g. param = dict(scale=0.05)
 
         Returns
         -------
-        gamma_1 : `float`
-            shear component in the x-direction
-        gamma_2 : `float`
-            shear component in the y-direction
-
-        Examples
-        --------
-        >>> from ler.lens_galaxy_population import LensGalaxyParameterDistribution
-        >>> lens = LensGalaxyParameterDistribution()
-        >>> gamma_1, gamma_2 = lens.shear_norm(size=1000)
+        lens_parameters : `dict`
+            dictionary of sampled lens parameters.
+            keys: sigma, q, phi, gamma, gamma1, gamma2
         """
 
-        if param:
-            scale = param["scale"]
+        lens_parameters = self.sampling_routine_sie_nsl(zl, zs, size=size)
 
-        if get_attribute:
-            return njit(lambda size: np.random.normal(loc=0, scale=scale,size=(2,size)))
-        else:
-            # Draw an external shear from a normal distribution
-            return np.random.normal(loc=0., scale=scale,size=(2,size))
+        # Sample the density profile slope distribution
+        gamma = self.density_profile_slope(size)
 
-    def mass_density_spectral_index_normal(
-        self, size=1000, mean=2.0, std=0.2, get_attribute=False, param=None
-    ):
+        # Sample shears
+        gamma1, gamma2 = self.external_shear.rvs(size)
+
+        # Add the lensing parameter dictionaries together
+        lens_parameters["gamma"] = gamma
+        lens_parameters["gamma1"] = gamma1
+        lens_parameters["gamma2"] = gamma2
+
+        return lens_parameters
+
+    def strongly_lensed_source_redshifts(self, size=1000):
         """
-        Function to sample the lens galaxy spectral index of the mass density profile from a normal distribution
+        Function to sample source redshifts, conditioned on the source being strongly lensed.
 
         Parameters
         ----------
         size : `int`
             number of lens parameters to sample
-        mean : `float`
-            mean of the normal distribution
-        std : `float`
-            standard deviation of the normal distribution
-        param : `dict`
-            Allows to pass in above parameters as dict.
-            e.g. param = dict(mean=2.0, std=0.2)
 
         Returns
         -------
-        gamma : `float`
-            spectral index of the density profile
+        redshifts : `float`
+            source redshifts conditioned on the source being strongly lensed
 
         Examples
         --------
         >>> from ler.lens_galaxy_population import LensGalaxyParameterDistribution
         >>> lens = LensGalaxyParameterDistribution()
-        >>> lens.mass_density_spectral_index_normal(size=1000)
-        """
-
-        if param:
-            mean = param["mean"]
-            std = param["std"]
-
-        if get_attribute:
-            return njit(lambda size: np.random.normal(loc=mean, scale=std, size=size))
-        else:
-            # Draw the spectral index from a normal distribution
-            return np.random.normal(loc=mean, scale=std, size=size)
-
-    def compute_einstein_radii(self, sigma, zl, zs):
+        >>> lens.strongly_lensed_source_redshifts(size=1000)
         """
-        Function to compute the Einstein radii of the lens galaxies
 
-        Parameters
-        ----------
-        sigma : `float`
-            velocity dispersion of the lens galaxy
-        zl : `float`
-            lens redshifts
-        zs : `float`
-            source redshifts
-
-        Returns
-        -------
-        theta_E : `float`
-            Einstein radii of the lens galaxies
+        z_max = self.z_max
 
-        Examples
-        --------
-        >>> from ler.lens_galaxy_population import LensGalaxyParameterDistribution
-        >>> lens = LensGalaxyParameterDistribution()
-        >>> sigma = 200.0
-        >>> zl = 0.5
-        >>> zs = 1.0
-        >>> lens.compute_einstein_radii(sigma, zl, zs)
-        """
+        def zs_function(zs_sl):
+            # get zs
+            # self.sample_source_redshifts from CBCSourceRedshiftDistribution class
+            zs = self.zs(size)  # this function is from CBCSourceParameterDistribution class
+            # put strong lensing condition with optical depth
+            tau = self.optical_depth(zs)
+            tau_max = self.optical_depth(np.array([z_max]))[0] # tau increases with z
+            r = np.random.uniform(0, tau_max, size=len(zs)) 
+            # Add the strongly lensed source redshifts to the list
+            # pick strongly lensed sources
+            zs_sl += list(zs[r < tau])  # list concatenation
 
-        # Compute the angular diameter distances
-        Ds = self.angular_diameter_distance(zs)
-        Dls = self.angular_diameter_distance_z1z2(zl, zs)
-        # Compute the Einstein radii
-        theta_E = (
-            4.0 * np.pi * (sigma / 299792.458) ** 2 * Dls / (Ds)
-        )  # Note: km/s for sigma; Dls, Ds are in Mpc
+            # Check if the zs_sl are larger than requested size
+            if len(zs_sl) >= size:
+                # Trim list to right size
+                zs_sl = zs_sl[:size]
+                return zs_sl
+            else:
+                # Run iteratively until we have the right number of lensing parmaeters
+                return zs_function(zs_sl)
 
-        return theta_E
+        zs_sl = []
 
-    def rjs_with_cross_section_SIS(self, param_dict):
+        return np.array(zs_function(zs_sl))
+        
+    def rjs_with_cross_section_sis(self, param_dict, cross_section_max=0.):
         """
-        Function to conduct rejection sampling wrt einstein radius
+            Function to conduct rejection sampling wrt einstein radius
 
-        Parameters
-        ----------
-        param_dict : `dict`
-            dictionary of lens parameters and source parameters
+            Parameters
+            ----------
+            param_dict : `dict`
+                dictionary of lens parameters and source parameters
 
-        Returns
-        -------
-        lens_params : `dict`
-            dictionary of lens parameters after rejection sampling
+            Returns
+            -------
+            lens_params : `dict`
+                dictionary of lens parameters after rejection sampling
         """
 
         theta_E = param_dict["theta_E"]
         size = len(theta_E)
-        theta_E_max = np.max(theta_E)  # maximum einstein radius
-        u = np.random.uniform(0, theta_E_max**2, size=size)
-        mask = u < theta_E**2
+        cross_section = np.pi * theta_E**2
+
+        max_ = np.max(cross_section)
+        if cross_section_max>max_:
+            max_ = cross_section_max
+
+        mask = np.random.uniform(size=size) < (cross_section / max_)
 
         # return the dictionary with the mask applied
-        return {key: val[mask] for key, val in param_dict.items()}
+        dict_ = {key: val[mask] for key, val in param_dict.items()}
+        dict_["cross_section"] = cross_section[mask]
+        return dict_, mask, max_
 
-    def rjs_with_cross_section_SIE(self, param_dict):
+    def rjs_with_cross_section_sie_feixu(self, param_dict, cross_section_max=0.):
         """
         Function to conduct rejection sampling wrt cross_section
 
@@ -926,260 +953,165 @@ class LensGalaxyParameterDistribution(CBCSourceParameterDistribution, ImagePrope
             dictionary of lens parameters after rejection sampling
         """
 
+        #print("rjs_with_cross_section_sie_feixu")
         theta_E = param_dict["theta_E"]
         q = param_dict["q"]
-        phi_cut = phi_cut_SIE(q)
         size = len(theta_E)
-        cross_section = theta_E**2 * phi_cut
-        max_ = np.max(cross_section)  # maximum einstein radius
+        cross_section = np.pi * theta_E**2 * phi_cut_SIE(q)
+
+        max_ = np.max(cross_section)
+        if cross_section_max>max_:
+            max_ = cross_section_max
+
         u = np.random.uniform(0, max_, size=size)
         mask = u < cross_section
 
         # return the dictionary with the mask applied
-        return {key: val[mask] for key, val in param_dict.items()}
-
-    @property
-    def sample_source_redshift_sl(self):
-        """
-        Function to sample source redshifts conditioned on the source being strongly lensed
-
-        Parameters
-        ----------
-        size : `int`
-            number samples to draw
-
-        Returns
-        -------
-        zs : `numpy.ndarray` (1D array of floats)
-            source redshifts conditioned on the source being strongly lensed
-
-        Examples
-        --------
-        >>> from ler.lens_galaxy_population import LensGalaxyParameterDistribution
-        >>> lens = LensGalaxyParameterDistribution()
-        >>> lens.sample_source_redshift_sl(size=1000)
-        """
-
-        return self._sample_source_redshift_sl
-
-    @sample_source_redshift_sl.setter
-    def sample_source_redshift_sl(self, prior):
-        try:
-            self._sample_source_redshift_sl = getattr(self, prior)
-        except:
-            self._sample_source_redshift_sl = prior
-
-    @property
-    def sample_source_parameters(self):
-        """
-        Function to sample source parameters conditioned on the source being strongly lensed
-
-        Parameters
-        ----------
-        size : `int`
-            number of lens parameters to sample
-
-        Returns
-        -------
-        source_parameters : `dict`
-            dictionary of source parameters conditioned on the source being strongly lensed
-        """
-        return self._sample_source_parameters
+        dict_ = {key: val[mask] for key, val in param_dict.items()}
+        dict_["cross_section"] = cross_section[mask]
+        return dict_, mask, max_
     
-    @sample_source_parameters.setter
-    def sample_source_parameters(self, prior):
-        try:
-            args = self.lens_param_samplers_params["source_parameters"]
-            self._sample_source_parameters = getattr(self, prior)(size=None, get_attribute=True, param=args)
-        except:
-            self._sample_source_parameters = prior
-
-    @property
-    def sample_lens_redshift(self):
-        """
-        Function to sample lens redshifts, conditioned on the lens being strongly lensed
-
-        Parameters
-        ----------
-        zs : `numpy.ndarray` (1D array of floats)
-            source redshifts
-
-        Returns
-        -------
-        zl : `numpy.ndarray` (1D array of floats)
-            lens redshifts corresponding to the source redshifts
-
-        Examples
-        --------
-        >>> from ler.lens_galaxy_population import LensGalaxyParameterDistribution
-        >>> lens = LensGalaxyParameterDistribution()
-        >>> zs = lens.sample_source_redshift_sl(size=1000)
-        >>> lens.sample_lens_redshift(zs=zs)
-        """
-
-        return self._sample_lens_redshift
-
-    @sample_lens_redshift.setter
-    def sample_lens_redshift(self, prior):
-        try:
-            args = self.lens_param_samplers_params["lens_redshift"]
-            self._sample_lens_redshift = getattr(self, prior)(zs=None, get_attribute=True, param=args)
-        except:
-            self._sample_lens_redshift = prior
-
-    @property
-    def sample_axis_rotation_angle(self):
+    def rjs_with_cross_section(self, param_dict, cross_section_max=0.):
         """
-        Function to sample the axis rotation angle of the elliptical lens galaxy from a uniform distribution
+            Function to conduct rejection sampling wrt cross_section of EPL+Shear lens
 
-        Parameters
-        ----------
-        size : `int`
-            number of lens parameters to sample
-
-        Returns
-        -------
-        phi : `float`
-            axis rotation angle of the elliptical lens galaxy
+            Parameters
+            ----------
+            param_dict : `dict`
+                dictionary of lens parameters and source parameters
 
-        Examples
-        --------
-        >>> from ler.lens_galaxy_population import LensGalaxyParameterDistribution
-        >>> lens = LensGalaxyParameterDistribution()
-        >>> lens.sample_axis_rotation_angle(size=1000)
+            Returns
+            -------
+            lens_params : `dict`
+                dictionary of lens parameters after rejection sampling
         """
-        return self._sample_axis_rotation_angle
 
-    @sample_axis_rotation_angle.setter
-    def sample_axis_rotation_angle(self, prior):
-        try:
-            args = self.lens_param_samplers_params["axis_rotation_angle"]
-            self._sample_axis_rotation_angle = getattr(self, prior)(size=None, get_attribute=True, param=args)
-        except:
-            self._sample_axis_rotation_angle = prior
+        theta_E_cut = 2.9243287409459857e-08 # this is numerically found
 
-    @property
-    def sample_shear(self):
-        """
-        Function to sample the elliptical lens galaxy shear from a normal distribution
-
-        Parameters
-        ----------
-        size : `int`
-            number of lens parameters to sample
+        # Pre-filter param_dict directly
+        size_original = len(param_dict["theta_E"])
+        idx_ = param_dict["theta_E"] > theta_E_cut
+        param_dict = {key: val[idx_]
+                    for key, val in param_dict.items()}
 
-        Returns
-        -------
-        gamma_1 : `float`
-            shear component in the x-direction
-        gamma_2 : `float`
-            shear component in the y-direction
-
-        Examples
-        --------
-        >>> from ler.lens_galaxy_population import LensGalaxyParameterDistribution
-        >>> lens = LensGalaxyParameterDistribution()
-        >>> gamma_1, gamma_2 = lens.shear_norm(size=1000)
-        """
+        size = len(param_dict["theta_E"])  # Update size after filtering
+        theta_E = param_dict["theta_E"]
+        e1, e2 = phi_q2_ellipticity_hemanta(param_dict["phi"], param_dict["q"])
+        gamma = param_dict["gamma"]
+        gamma1 = param_dict["gamma1"]
+        gamma2 = param_dict["gamma2"]
 
-        return self._sample_shear
+        cross_section = self.cross_section(theta_E, e1, e2, gamma, gamma1, gamma2)
 
-    @sample_shear.setter
-    def sample_shear(self, prior):
-        try:
-            args = self.lens_param_samplers_params["shear"]
-            self._sample_shear = getattr(self, prior)(size=None, get_attribute=True, param=args)
-        except:
-            self._sample_shear = prior
-
-    @property
-    def sample_mass_density_spectral_index(self):
-        """
-        Function to sample the lens galaxy spectral index of the mass density profile from a normal distribution
+        max_ = np.max(cross_section)
+        if cross_section_max>max_:
+            max_ = cross_section_max
 
-        Parameters
-        ----------
-        size : `int`
-            number of lens parameters to sample
-
-        Returns
-        -------
-        gamma : `float`
-            spectral index of the density profile
-
-        Examples
-        --------
-        >>> from ler.lens_galaxy_population import LensGalaxyParameterDistribution
-        >>> lens = LensGalaxyParameterDistribution()
-        >>> lens.mass_density_spectral_index_normal(size=1000)
-        """
+        u = np.random.uniform(0, max_, size=size)
+        mask = u < cross_section
 
-        return self._sample_mass_density_spectral_index
+        # return the dictionary with the mask applied
+        dict_ = {key: val[mask] for key, val in param_dict.items()}
+        dict_["cross_section"] = cross_section[mask]
 
-    @sample_mass_density_spectral_index.setter
-    def sample_mass_density_spectral_index(self, prior):
-        try:
-            args = self.lens_param_samplers_params["mass_density_spectral_index"]
-            self._sample_mass_density_spectral_index = getattr(self, prior)(size=None, get_attribute=True, param=args)
-        except:
-            self._sample_mass_density_spectral_index = prior
+        mask_complete = np.zeros(size_original, dtype=bool)
+        mask_complete[idx_] = mask
+        return dict_, mask_complete, max_
 
-    @property
-    def sample_source_parameters(self):
+    
+    def rjs_with_cross_section_mp(self, param_dict, cross_section_max=0.):
         """
-        Function to sample source parameters conditioned on the source being strongly lensed
+            Function to conduct rejection sampling wrt cross_section, multiprocessing
 
-        Parameters
-        ----------
-        size : `int`
-            number of lens parameters to sample
+            Parameters
+            ----------
+            param_dict : `dict`
+                dictionary of lens parameters and source parameters
 
-        Returns
-        -------
-        source_parameters : `dict`
-            dictionary of source parameters conditioned on the source being strongly lensed
+            Returns
+            -------
+            lens_params : `dict`
+                dictionary of lens parameters after rejection sampling
         """
-        return self._sample_source_parameters
-
-    @sample_source_parameters.setter
-    def sample_source_parameters(self, prior):
-        try:
-            self._sample_source_parameters = getattr(self, prior)
-        except:
-            self._sample_source_parameters = prior
 
+        theta_E_cut = 1e-09 # this is numerically found
+        theta_E = param_dict["theta_E"]
+        size_original = len(theta_E)
 
-    @property
-    def available_lens_prior_list_and_its_params(self):
-        """
-        Dictionary with list all the available priors and it's corresponding parameters. This is an immutable instance attribute.
-        """
+        # Pre-filter param_dict directly
         
-        self._available_lens_prior_list_and_its_params = dict(
-            source_redshift_sl=dict(strongly_lensed_source_redshifts=None),
-            lens_redshift=dict(lens_redshift_SDSS_catalogue=None),
-            velocity_dispersion=self.available_velocity_dispersion_list_and_its_params,
-            axis_ratio=self.available_axis_ratio_list_and_its_params,
-            axis_rotation_angle=dict(axis_rotation_angle_uniform=dict(phi_min=0.0, phi_max=2 * np.pi)),
-            shear=dict(shear_norm=dict(scale=0.05)),
-            mass_density_spectral_index=dict(mass_density_spectral_index_normal=dict(mean=2.0, std=0.2)),
-            source_parameters=dict(sample_gw_parameters=None),
-        )
-
-        return self._available_lens_prior_list_and_its_params
-    
-    @property
-    def available_lens_functions(self):
-        """
-        Dictionary with list all the available lens functions. This is an immutable instance attribute.
-        """
-
-        self._available_lens_functions = dict(
-            strong_lensing_condition=["rjs_with_cross_section_SIE", "rjs_with_cross_section_SIS"],
-            optical_depth=["SIS", "optical_depth_SIS_haris","optical_depth_SIS_hemanta", "SIE", "optical_depth_SIE_hemanta"],
-            param_sampler_type=["sample_all_routine"],
-        )
-
-        return self._available_lens_functions
-
+        idx_ = theta_E > theta_E_cut
+        if np.sum(idx_) > 0:
+        
+            param_dict = {key: val[idx_] 
+                        for key, val in param_dict.items()}
+            size = np.sum(idx_)
+
+            e1, e2 = phi_q2_ellipticity_hemanta(param_dict["phi"], param_dict["q"])
+
+            # Combine parameters into a single array for multiprocessing
+            params = np.array([
+                param_dict["theta_E"],
+                e1,
+                e2,
+                param_dict["gamma"],
+                param_dict["gamma1"],
+                param_dict["gamma2"],
+                np.arange(size, dtype=int)
+            ]).T
+
+            cross_section = np.zeros(size)  # Directly create filtered array
+            # with Pool(processes=self.npool) as pool:
+            #     for result in tqdm(pool.imap_unordered(cross_section_mp, params), total=size,
+            #                     ncols=100):  # Remove total from tqdm
+            #         idx_, tau_ = result
+            #         cross_section[idx_] = tau_
+            # without tqdm. Use map
+            # Use multiprocessing to calculate cross-section values
+            with Pool(processes=self.npool) as pool:
+                for idx, tau in pool.imap_unordered(cross_section_mp, params):
+                    cross_section[idx] = tau
+
+            # Perform rejection sampling
+            max_ = np.max(cross_section)
+
+            if cross_section_max>max_:
+                max_ = cross_section_max
+            random_values = np.random.uniform(0, max_, size=size)
+            mask = random_values < cross_section
+        else:
+            print("No valid values found")
+            mask = np.zeros(size, dtype=bool)
+
+        # Return the dictionary with the mask applied
+        dict_ = {key: val[mask] for key, val in param_dict.items()}
+        dict_["cross_section"] = cross_section[mask]
+
+        mask_complete = np.zeros(size_original, dtype=bool)
+        mask_complete[idx_] = mask
+        return dict_, mask_complete, max_
+
+    # @property
+    # def sample_source_parameters(self):
+    #     """
+    #     Function to sample source parameters conditioned on the source being strongly lensed
+
+    #     Parameters
+    #     ----------
+    #     size : `int`
+    #         number of lens parameters to sample
+
+    #     Returns
+    #     -------
+    #     source_parameters : `dict`
+    #         dictionary of source parameters conditioned on the source being strongly lensed
+    #     """
+    #     return self._sample_source_parameters
+
+    # @sample_source_parameters.setter
+    # def sample_source_parameters(self, prior):
+    #     try:
+    #         self._sample_source_parameters = getattr(self, prior)
+    #     except:
+    #         self._sample_source_parameters = prior