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

ler/image_properties/image_properties.py

@@ -120,8 +120,7 @@ class ImageProperties():
                  z_max=10,
                  n_min_images=2, 
                  n_max_images=4,
-                 geocent_time_min=1126259462.4,
-                 geocent_time_max=1126259462.4+365*24*3600*20,
+                 time_window=365*24*3600*20,
                  lens_model_list=['EPL_NUMBA', 'SHEAR'],
                  cosmology=None,
                  spin_zero=True,
@@ -136,8 +135,9 @@ class ImageProperties():
         self.lens_model_list = lens_model_list  # list of lens models
         self.spin_zero = spin_zero
         self.spin_precession = spin_precession
-        self.geocent_time_min = geocent_time_min
-        self.geocent_time_max = geocent_time_max
+        self.time_window = time_window
+        # self.geocent_time_min = geocent_time_min
+        # self.geocent_time_max = geocent_time_max
         self.cosmo = cosmology if cosmology else cosmo
         
         # initialize the interpolator's parameters
@@ -368,15 +368,20 @@ class ImageProperties():
         -------
             snrs : `dict`
                 signal to noise ratio for each image in each event.
-                (dictionary containing 'H1', 'L1', ..., and 'optimal_snr_net', which is the network snr, for each image as an array with dimensions (number_of_lensed_events,n_max_images) )
+                (dictionary containing 'H1', 'L1', ..., and 'snr_net', which is the network snr, for each image as an array with dimensions (number_of_lensed_events,n_max_images) )
 
         """
         # needed to calculate effective luminosity distance and effective time delay
         n_max_images = self.n_max_images
         magnifications = lensed_param["magnifications"]
         time_delays = lensed_param["time_delays"]
+        imgage_type = lensed_param["image_type"]
         size = len(magnifications)
 
+        # image type to morse phase
+        imgage_type[imgage_type==1.] = 0.
+        imgage_type[imgage_type==2.] = np.pi/2
+
         # Get the binary parameters
         number_of_lensed_events = len(magnifications)
         mass_1, mass_2, theta_jn, psi, ra, dec, phase, a_1, a_2, tilt_1, tilt_2, phi_12, phi_jl  = (
@@ -407,7 +412,7 @@ class ImageProperties():
         # setting up snr dictionary
         result_dict = dict()
         if snr_calculator:
-            result_dict["optimal_snr_net"] = (
+            result_dict["snr_net"] = (
                 np.ones((number_of_lensed_events, n_max_images)) * np.nan
             )
         # setting up pdet dictionary
@@ -445,9 +450,23 @@ class ImageProperties():
             time_eff_present = True
         else:
             raise ValueError("geocent_time or effective_geocent_time not given")
+        
+        if "phase" in lensed_param:
+            phase = lensed_param["phase"]
+            lensed_param["effective_phase"] = np.ones((number_of_lensed_events, n_max_images)) * np.nan
+            phase_eff_present = False
+        elif "effective_phase" in lensed_param:
+            phase_eff = lensed_param["effective_phase"]
+            phase_eff_present = True
+        else:
+            raise ValueError("phase or effective_phase not given")
+
 
         # Get the optimal signal to noise ratios for each image
         # iterate over the image type (column)
+        geocent_time_min = np.min(geocent_time)
+        geocent_time_max = geocent_time_min + self.time_window
+
         for i in range(n_max_images):
             
             # get the effective time for each image type
@@ -455,8 +474,10 @@ class ImageProperties():
                 effective_geocent_time = geocent_time + time_delays[:, i]
             else:
                 effective_geocent_time = time_eff[:, i]
+                
             # choose only the events that are within the time range and also not nan
-            idx = (effective_geocent_time <= self.geocent_time_max) & (effective_geocent_time >= self.geocent_time_min)
+            idx = (effective_geocent_time <= geocent_time_max) & (effective_geocent_time >= geocent_time_min)
+
             # get the effective luminosity distance for each image type
             if not dl_eff_present:
                 effective_luminosity_distance = luminosity_distance / np.sqrt(
@@ -464,12 +485,17 @@ class ImageProperties():
                 )
             else:
                 effective_luminosity_distance = dl_eff[:, i]
+            # get the effective phase for each image type
+            if not phase_eff_present:
+                effective_phase = phase - imgage_type[:, i]  # morse phase correction
+            else:
+                effective_phase = phase_eff[:, i]
             
             # check for nan values
-            idx = idx & ~np.isnan(effective_luminosity_distance) & ~np.isnan(effective_geocent_time)
+            idx = idx & ~np.isnan(effective_luminosity_distance) & ~np.isnan(effective_geocent_time) & ~np.isnan(effective_phase)
 
             # Each image has their own effective luminosity distance and effective geocent time
-            if len(effective_luminosity_distance) != 0:
+            if sum(idx) != 0:
                 # Returns a dictionary
                 if snr_calculator:
                     optimal_snr = snr_calculator(
@@ -479,7 +505,7 @@ class ImageProperties():
                             luminosity_distance=effective_luminosity_distance[idx],
                             theta_jn=theta_jn[idx],
                             psi=psi[idx],
-                            phase=phase[idx],
+                            phase= effective_phase[idx],
                             geocent_time=effective_geocent_time[idx],
                             ra=ra[idx],
                             dec=dec[idx],
@@ -492,7 +518,7 @@ class ImageProperties():
                         ),
                         output_jsonfile=False,
                     )
-                    result_dict["optimal_snr_net"][idx, i] = optimal_snr["optimal_snr_net"]
+                    result_dict["snr_net"][idx, i] = optimal_snr["snr_net"]
 
                     if list_of_detectors:
                         for detector in list_of_detectors:
@@ -507,7 +533,7 @@ class ImageProperties():
                             luminosity_distance=effective_luminosity_distance[idx],
                             theta_jn=theta_jn[idx],
                             psi=psi[idx],
-                            phase=phase[idx],
+                            phase= effective_phase[idx],
                             geocent_time=effective_geocent_time[idx],
                             ra=ra[idx],
                             dec=dec[idx],
@@ -529,10 +555,13 @@ class ImageProperties():
 
                 lensed_param["effective_luminosity_distance"][:, i] = effective_luminosity_distance
                 lensed_param["effective_geocent_time"][:, i] = effective_geocent_time
+                lensed_param["effective_phase"][:, i] = effective_phase
 
         if dl_eff_present:
             del lensed_param["effective_luminosity_distance"]
         if time_eff_present:
             del lensed_param["effective_geocent_time"]
+        if phase_eff_present:
+            del lensed_param["effective_phase"]
 
         return result_dict, lensed_param

ler/image_properties/multiprocessing_routine.py

@@ -7,8 +7,8 @@ import numpy as np
 from lenstronomy.LensModel.lens_model import LensModel
 from lenstronomy.LensModel.Solver.lens_equation_solver import LensEquationSolver
 from lenstronomy.LensModel.Solver.epl_shear_solver import caustics_epl_shear
-from lenstronomy.Util.param_util import phi_q2_ellipticity
-from ..lens_galaxy_population.jit_functions import axis_ratio_rayleigh
+# from lenstronomy.Util.param_util import phi_q2_ellipticity
+# from ..lens_galaxy_population.jit_functions import axis_ratio_rayleigh
 
 # For sampling from caustic
 from shapely.geometry import Polygon
@@ -38,7 +38,7 @@ def solve_lens_equation(lens_parameters):
     Returns
     -------
     x_source : `float`
-        x position of the source in the source plane
+        x position of the source in the source plane, unit: arcsec
     y_source : `float`
         y position of the source in the source plane
     x0_image_position : `float`
@@ -121,7 +121,7 @@ def solve_lens_equation(lens_parameters):
         )
         caustic = np.logical_not(np.isnan(caustic_double_points).any())
 
-        # If there is a nan, caustic=False, draw a new gamma
+        # If there is a nan, caustic=False, this batch will be ignored and resampled outside in lens_parameters sampling
         if caustic:
             break
         else:
@@ -139,6 +139,7 @@ def solve_lens_equation(lens_parameters):
                 iteration,
             )
 
+    # define region in the source plane where 2 or more images are formed
     caustic_double = Polygon(caustic_double_points.T)
 
     # check for strong lensed condition
@@ -242,208 +243,3 @@ def solve_lens_equation(lens_parameters):
     )
 
 
-# def solve_lens_equation(lens_parameters):
-#     """
-#     Function to solve the lens equation (min_image = 2)
-
-#     Parameters
-#     ----------
-#     lens_parameters : `list`
-#         a list of parameters
-#         lens_parameters[0] = min_images : minimum number of images
-#         lens_parameters[1] = e1 : ellipticity
-#         lens_parameters[2] = e2 : ellipticity
-#         lens_parameters[3] = gamma : power-law index
-#         lens_parameters[4] = gamma1 : shear
-#         lens_parameters[5] = gamma2 : shear
-#         lens_parameters[6] = zl : redshift of the lens
-#         lens_parameters[7] = zs : redshift of the source
-#         lens_parameters[8] = einstein_radius : Einstein radius
-#         lens_parameters[9] = iteration : iteration number
-#         lens_parameters[10:] = lens_model_list : numpy array of lens models
-
-#     Returns
-#     -------
-#     x_source : `float`
-#         x position of the source in the source plane
-#     y_source : `float`
-#         y position of the source in the source plane
-#     x0_image_position : `float`
-#         x position of the images in the source plane
-#     x1_image_position : `float`
-#         y position of the images in the source plane
-#     magnifications : `float`
-#         magnification of the images
-#     time_delays : `float`
-#         time-delay of the images
-#     nImages : `int`
-#         number of images
-#     determinant : `float`
-#         determinant of the hessian matrix
-#     trace : `float`
-#         trace of the hessian matrix
-#     iteration : `int`
-#         iteration number
-
-#     Examples
-#     --------
-#     >>> from ler.multiprocessing_routine import solve_lens_equation1
-#     >>> import numpy as np
-#     >>> from multiprocessing import Pool
-#     >>> # lens parameters input contains 12 parameters [e1, e2, gamma, gamma1, gamma2, zl, zs, einstein_radius, iteration, lens_model_list]
-#     >>> lens_parameters1 = np.array([2, 0.024069457093642648, -0.016002190961948142, 1.8945414936459974, 0.10117465203892329, 0.09600089396968613, 0.2503743800068136, 0.9418211055453296, 2.5055790287104725e-06, 0, 'EPL_NUMBA', 'SHEAR'], dtype=object)
-#     >>> lens_parameters2 = np.array([2, -0.04030088581646998, -0.01419438113690042, 2.0068239327017, 0.08482718989370612, -0.015393332086560785, 1.0952303138971118, 2.5534097159384417, 1.0125570159563301e-06, 1, 'EPL_NUMBA', 'SHEAR'], dtype=object)
-#     >>> input_arguments = np.vstack((lens_parameters1, lens_parameters2))
-#     >>> # solve the lens equation for each set of lens parameters
-#     >>> with Pool(2) as p:
-#     ...     result = p.map(solve_lens_equation1, input_arguments)
-#     >>> # result is a list of tuples
-#     >>> # each tuple contains the output parameters of the function
-#     >>> # each output parameter contains x_source, y_source, x0_image_position, x1_image_position, magnifications, time_delays, nImages, determinant, trace, iteration
-#     >>> print(f"magnification of images with lens parameters 'lens_parameters1' is {result[0][6]}")
-#     magnification of images with lens parameters 'lens_parameters1' is [ 2.18973765 -1.27542831]
-
-#     """
-#     n_min_images = int(lens_parameters[0])
-#     zl = lens_parameters[6]
-#     zs = lens_parameters[7]
-#     einstein_radius = lens_parameters[8]
-#     iteration = lens_parameters[9]
-#     # lensModel parameters are the same for the three functions used for image param calculation
-#     # 1. x-y position of images in the source plane, 2. magnifications, 3. time-delays (relative)
-#     lensModel = LensModel(
-#         lens_model_list=lens_parameters[10:].tolist(), z_lens=zl, z_source=zs
-#     )
-
-#     lens_eq_solver = LensEquationSolver(lensModel)
-
-#     factor = 1.0
-#     # ---------------------------------------------------#
-#     #     x-y position of images in the source plane
-#     # ---------------------------------------------------#
-#     # Get the caustic curve cut by the lens
-#     # First check if there is any nan in the caustic points
-#     def check_caustic():
-#         """
-#         Function to check if there is any nan in the caustic points
-#         """
-#         while True:
-#             kwargs_lens = [
-#                 {
-#                     "theta_E": factor,
-#                     "e1": lens_parameters[1],
-#                     "e2": lens_parameters[2],
-#                     "gamma": lens_parameters[3],
-#                     "center_x": 0.0,
-#                     "center_y": 0.0,
-#                 },
-#                 {
-#                     "gamma1": lens_parameters[4],
-#                     "gamma2": lens_parameters[5],
-#                     "ra_0": 0,
-#                     "dec_0": 0,
-#                 },
-#             ]
-#             caustic_double_points = caustics_epl_shear(
-#                 kwargs_lens, return_which="double", maginf=-100
-#             )
-#             caustic = np.logical_not(np.isnan(caustic_double_points).any())
-
-#             # If there is a nan, caustic=False, draw a new gamma
-#             if caustic:
-#                 break
-#             else:
-#                 #
-#                 q = axis_ratio_rayleigh(sigma=np.array([160.]))[0]
-#                 phi = np.random.uniform(0.0, 2*np.pi, size=1)[0]
-#                 lens_parameters[1], lens_parameters[2] = phi_q2_ellipticity(phi, q)
-#                 # resampling power-law index and shear
-#                 lens_parameters[3] = np.random.normal(loc=2.0, scale=0.2, size=1)[0]
-#                 gamma1, gamma2 = np.random.normal(loc=0, scale=0.05,size=(2,1))
-#                 lens_parameters[4], lens_parameters[5] = gamma1[0], gamma2[0]
-
-#         caustic_double = Polygon(caustic_double_points.T)
-#         return caustic_double, kwargs_lens
-    
-#     caustic_double, kwargs_lens = check_caustic()
-
-#     # check for strong lensed condition
-#     strongly_lensed = False
-#     i = 0
-#     while strongly_lensed == False:
-#         # Draw random points within the caustic
-#         # sometimes x_source, y_source positions are at same location and the solver fails
-#         # so we use a try-except block to catch the error and draw a new point
-#         # try:
-#         x_source, y_source = pointpats.random.poisson(caustic_double, size=1)
-#         # Solve the lens equation
-#         (
-#             x0_image_position,
-#             x1_image_position,
-#         ) = lens_eq_solver.image_position_from_source(
-#             sourcePos_x=x_source,
-#             sourcePos_y=y_source,
-#             kwargs_lens=kwargs_lens,
-#             solver="analytical",
-#             magnification_limit=1.0 / 1000.0,
-#             arrival_time_sort=True,
-#         )
-#         nImages = len(x0_image_position)  # shows how many images
-#         if nImages >= n_min_images:
-#             strongly_lensed = True
-#         # except:
-#         #     pass
-#         ## test ##
-#         if i > 10:
-#             print("Could not find a valid source position. Run the sampling again by setting resume=True")
-#             caustic_double, kwargs_lens = check_caustic()
-#             i = 0
-#         i += 1
-#         ##########
-
-#     # ---------------------------------------------------#
-#     #          magnification and time-delay
-#     # ---------------------------------------------------#
-#     # theta_E is in arcsec
-#     theta_E_nImages = einstein_radius * np.ones(nImages)
-#     radian_to_arcseconds = 180.0 / np.pi * 3600.0
-#     days_to_seconds = 24.0 * 3600.0
-#     # can have multiple magnification
-#     magnifications = lensModel.magnification(
-#         x0_image_position, x1_image_position, kwargs_lens
-#     )
-#     time_delays = (
-#         lensModel.arrival_time(x0_image_position, x1_image_position, kwargs_lens)
-#         * (theta_E_nImages * radian_to_arcseconds) ** 2
-#         * days_to_seconds
-#     )
-
-#     # ---------------------------------------------------#
-#     #     Params needed for image-type classification
-#     # ---------------------------------------------------#
-#     # it is faster to use numpy array operation to do image classification
-#     # return: f_xx, f_xy, f_yx, f_yy components
-#     hessian = lensModel.hessian(x0_image_position, x1_image_position, kwargs_lens)
-#     determinant = np.array(
-#         (1 - hessian[0]) * (1 - hessian[3]) - hessian[1] * hessian[2]
-#     )
-#     trace = np.array(2 - hessian[0] - hessian[3])
-
-#     #  return also gamma1, gamma2
-#     return (
-#         x_source,
-#         y_source,
-#         x0_image_position,
-#         x1_image_position,
-#         magnifications,
-#         time_delays,
-#         nImages,
-#         determinant,
-#         trace,
-#         lens_parameters[1],
-#         lens_parameters[2],
-#         lens_parameters[3],
-#         lens_parameters[4],
-#         lens_parameters[5],
-#         iteration,
-#     )

ler/lens_galaxy_population/__init__.py

@@ -2,3 +2,5 @@ from .lens_galaxy_parameter_distribution import *
 from .optical_depth import *
 from .mp import *
 from .jit_functions import *
+
+from .import lens_param_data

ler/lens_galaxy_population/jit_functions.py

@@ -68,9 +68,10 @@ def pdf_phi_z_div_0(s, z):
     return phi_sim_z / phi_sim_0
 
 @njit
-def phi(s,z, cosmology_h=0.7):
+def phi(s, z, alpha=0.94, beta=1.85, phistar=2.099e-2, sigmastar=113.78):
     """
     Function to calculate the lens galaxy velocity dispersion function at redshift z.
+    For Oguri et al. (2018b) model: alpha=0.94, beta=1.85, phistar=2.099e-2*(self.cosmo.h/0.7)**3, sigmastar=113.78
 
     Parameters
     ----------
@@ -85,15 +86,17 @@ def phi(s,z, cosmology_h=0.7):
     -------
     result : `float: array`
     """
-
-    result = s**4*pdf_phi_z_div_0(s,z)*phi_loc_bernardi(sigma=s, cosmology_h=cosmology_h)
-    # result[result < 0.] = 0.
+    
+    result = pdf_phi_z_div_0(s,z) * phi_loc_bernardi(sigma=s, alpha=alpha, beta=beta, phistar=phistar, sigmastar=sigmastar)
+    result[result < 0.] = 0.
     return result
 
 @njit
-def phi_loc_bernardi(sigma, alpha=0.94, beta=1.85, phistar=2.099e-2, sigmastar=113.78, cosmology_h=0.7):
+def phi_loc_bernardi(sigma, alpha, beta, phistar, sigmastar):
     """
     Function to calculate the local universe velocity dispersion function. Bernardi et al. (2010).
+    For Oguri et al. (2018b) model: alpha=0.94, beta=1.85, phistar=2.099e-2*(self.cosmo.h/0.7)**3, sigmastar=113.78
+    For Choi et al. (2008) model: alpha = 2.32 / 2.67, beta = 2.67, phistar = 8.0e-3*self.cosmo.h**3, sigmastar = 161.0
 
     Parameters
     ----------
@@ -109,7 +112,7 @@ def phi_loc_bernardi(sigma, alpha=0.94, beta=1.85, phistar=2.099e-2, sigmastar=1
     philoc_ : `float: array`
     """
 
-    phistar = phistar * (cosmology_h / 0.7) ** 3  # Mpc**-3
+    # phistar = phistar * (cosmology_h / 0.7) ** 3  # Mpc**-3
     philoc_ = phistar*(sigma/sigmastar)**alpha * np.exp(-(sigma/sigmastar)**beta) * beta/gamma_(alpha/beta)/sigma
     return philoc_
 
@@ -143,7 +146,7 @@ def phi_cut_SIE(q):
     return result/np.pi
 
 @njit
-def axis_ratio_rayleigh(sigma, q_min=0.2, q_max=1.0):
+def axis_ratio_rayleigh_rvs(sigma, q_min=0.2, q_max=1.0):
         """
         Function to sample axis ratio from rayleigh distribution with given velocity dispersion.
 
@@ -184,6 +187,16 @@ def axis_ratio_rayleigh(sigma, q_min=0.2, q_max=1.0):
 
         return q
 
+@njit
+def axis_ratio_rayleigh_pdf(q, sigma, q_min=0.2, q_max=1.0):
+
+    s = 0.38 - 5.7e-4 * sigma
+    idx = (q >= q_min) & (q <= q_max)
+    q_pdf = np.zeros_like(q)
+    q_pdf[idx] = (1-q[idx])/s[idx]**2 * np.exp(-0.5*(1-q[idx])**2/s[idx]**2)
+    return q_pdf
+
+
 @njit
 def velocity_dispersion_z_dependent(size, zl, zl_list, vd_inv_cdf):
     """
@@ -215,9 +228,9 @@ def velocity_dispersion_z_dependent(size, zl, zl_list, vd_inv_cdf):
     return samples
 
 @njit
-def lens_redshift_SDSS_catalogue(zs, splineDc, splineDcInv, u, cdf):
+def lens_redshift_SDSS_catalogue_sis(zs, splineDc, splineDcInv, u, cdf):
     """
-    Function to sample lens redshift from the SDSS catalogue.
+    Function to sample lens redshift from the SDSS catalogue. Haris et al. (2018) cdf = (10 * u**3 - 15 * u**4 + 6 * u**5)
 
     Parameters
     ----------
@@ -275,4 +288,83 @@ def bounded_normal_sample(size, mean, std, low, high):
         samples[i] = sample
     return samples
 
+
+@njit
+def phi_q2_ellipticity_hemanta(phi, q):
+    """Function to convert phi and q to ellipticity e1 and e2.
+
+    Parameters
+    ----------
+    phi : `float: array`
+        angle of the major axis in radians
+    q : `float: array`
+        axis ratio
+
+    Returns
+    -------
+    e1 : `float: array`
+        ellipticity component 1
+    e2 : `float: array`
+    """
+
+    e_1 = (1.0 - q) / (1.0 + q) * np.cos(2 * phi)
+    e_2 = (1.0 - q) / (1.0 + q) * np.sin(2 * phi)
+    return e_1, e_2
+
+# @njit
+# def sample_sigma_zl(pdf, sigma_min, sigma_max, zl_min, zl_max, zs, chunk_size=10000):
+#     x_sample = []
+#     y_sample = []
     
+#     for z in zs:
+#         xmin, xmax = sigma_min, sigma_max
+#         ymin, ymax = zl_min, z
+        
+#         # Keep sampling until a valid (sigma, zl) pair is found for this z
+#         while True:
+#             x_try = np.random.uniform(xmin, xmax, chunk_size)
+#             y_try = np.random.uniform(ymin, ymax, chunk_size)
+#             pdf_xy_try = pdf(x_try, y_try)
+#             zmax = np.max(pdf_xy_try)  # Maximum of the PDF for the current batch
+
+#             # Generate acceptance thresholds
+#             z_try = np.random.uniform(0, zmax, chunk_size)
+
+#             # Check which samples are accepted
+#             accepted_indices = z_try < pdf_xy_try
+#             if np.any(accepted_indices):
+#                 # Accept the first valid sample
+#                 x_sample.append(x_try[accepted_indices][0])
+#                 y_sample.append(y_try[accepted_indices][0])
+#                 break  # Exit the loop once the first valid sample is collected
+
+#     return np.array(x_sample), np.array(y_sample)
+
+@njit
+def sample_sigma_zl(pdf, sigma_min, sigma_max, zl_min, zl_max, zs, chunk_size=10000):
+    x_sample = []
+    y_sample = []
+    
+    for z in zs:
+        xmin, xmax = sigma_min, sigma_max
+        ymin, ymax = zl_min, zl_max  # Use full range and then filter
+        
+        valid_sample_found = False
+        while not valid_sample_found:
+            x_try = np.random.uniform(xmin, xmax, chunk_size)
+            y_try = np.random.uniform(ymin, ymax, chunk_size)
+            pdf_xy_try = pdf(x_try, y_try)
+            zmax = np.max(pdf_xy_try)  # Find maximum PDF value for current batch
+
+            # Generate acceptance thresholds
+            z_try = np.random.uniform(0, zmax, chunk_size)
+
+            # Check and accept samples
+            accepted = (z_try < pdf_xy_try) & (y_try < z)  # Ensure zl < zs condition
+            if np.any(accepted):
+                idx = np.argmax(accepted)  # First valid index
+                x_sample.append(x_try[idx])
+                y_sample.append(y_try[idx])
+                valid_sample_found = True
+
+    return np.array(x_sample), np.array(y_sample)