ler 0.3.9__tar.gz → 0.4.1__tar.gz

{ler-0.3.9 → ler-0.4.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ler
-Version: 0.3.9
+Version: 0.4.1
 Summary: Gravitational waves Lensing Rates
 Home-page: https://github.com/hemantaph/ler
 Author: Hemantakumar

ler-0.4.1/ler/__init__.py

@@ -0,0 +1,68 @@
+"""
+LeR
+"""
+
+# mypackage/cli.py
+import argparse
+# import subprocess, os, sys, signal, warnings
+
+## import pycbc
+import os
+import multiprocessing as mp
+
+def set_multiprocessing_start_method():
+    if os.name == 'posix':  # posix indicates the program is run on Unix/Linux/Mac
+      print("Setting multiprocessing start method to 'fork'")
+      try:
+         mp.set_start_method('fork', force=True)
+      except RuntimeError:
+         # The start method can only be set once and must be set before any process starts
+         pass
+    else:
+      print("Setting multiprocessing start method to 'spawn'")
+      # For Windows and other operating systems, use 'spawn'
+      try:
+         mp.set_start_method('spawn', force=True)
+      except RuntimeError:
+         pass
+
+set_multiprocessing_start_method()
+
+# try:
+#    mp.set_start_method('fork', force=True)
+# except RuntimeError:
+#    pass
+
+# if sys.platform == 'darwin':
+#     HAVE_OMP = False
+
+#     # MacosX after python3.7 switched to 'spawn', however, this does not
+#     # preserve common state information which we have relied on when using
+#     # multiprocessing based pools.
+#     import multiprocessing
+#     if multiprocessing.get_start_method(allow_none=True) is None:
+#         if hasattr(multiprocessing, 'set_start_method'):
+#             multiprocessing.set_start_method('fork')
+#     elif multiprocessing.get_start_method() != 'fork':
+#         warnings.warn("PyCBC requires the use of the 'fork' start method for multiprocessing, it is currently set to {}".format(multiprocessing.get_start_method()))
+# else:
+#     HAVE_OMP = True
+
+__author__ = 'hemanta_ph <hemantaphurailatpam@gmail.com>'
+
+# The version as used in the setup.py
+__version__ = "0.4.1"
+
+# add __file__
+import os
+__file__ = os.path.abspath(__file__)
+
+from . import rates, gw_source_population, lens_galaxy_population, image_properties, utils
+
+from .rates import ler, gwrates
+from .gw_source_population import cbc_source_parameter_distribution, cbc_source_redshift_distribution
+from .lens_galaxy_population import lens_galaxy_parameter_distribution, optical_depth
+from .image_properties import image_properties
+from .utils import utils, plots
+
+

{ler-0.3.9 → ler-0.4.1}/ler/gw_source_population/cbc_source_parameter_distribution.py

@@ -58,8 +58,8 @@ class CBCSourceParameterDistribution(CBCSourceRedshiftDistribution):
         If True, spin parameters are completely ignore in the sampling.
         default: True
     spin_precession : `bool`
-        If spin_zero=True and spin_precession=True, spin parameters are sampled for precessing binaries.
-        if spin_zero=True and spin_precession=False, spin parameters are sampled for aligned/anti-aligned spin binaries.
+        If spin_zero=False and spin_precession=True, spin parameters are sampled for precessing binaries.
+        if spin_zero=False and spin_precession=False, spin parameters are sampled for aligned/anti-aligned spin binaries.
         default: False
     directory : `str`
         Directory to store the interpolator pickle files

{ler-0.3.9 → ler-0.4.1}/ler/image_properties/image_properties.py

@@ -6,6 +6,8 @@ The class inherits from the CompactBinaryPopulation class, which is used to samp
 
 import warnings
 warnings.filterwarnings("ignore")
+import logging
+logging.getLogger('numexpr.utils').setLevel(logging.ERROR)
 # for multiprocessing
 from multiprocessing import Pool
 from tqdm import tqdm

{ler-0.3.9 → ler-0.4.1}/ler/lens_galaxy_population/lens_galaxy_parameter_distribution.py

@@ -374,6 +374,8 @@ class LensGalaxyParameterDistribution(CBCSourceParameterDistribution, ImagePrope
             lens_model_list=["EPL_NUMBA", "SHEAR"],
         )
         input_params_image.update(params)
+
+        print("input_params_image", input_params_image)
         ImageProperties.__init__(
             self,
             npool=self.npool,

{ler-0.3.9 → ler-0.4.1}/ler/rates/ler.py

@@ -6,14 +6,22 @@ This module contains the main class for calculating the rates of detectable grav
 import os
 import warnings
 warnings.filterwarnings("ignore")
+import logging
+logging.getLogger('numexpr.utils').setLevel(logging.ERROR)
 import contextlib
 import numpy as np
 from scipy.stats import norm
 from astropy.cosmology import LambdaCDM
 from ..lens_galaxy_population import LensGalaxyParameterDistribution
-from ..utils import load_json, append_json, get_param_from_json, batch_handler, add_dict_values
+from ..utils import load_json, append_json, get_param_from_json, batch_handler
 
 
+# # multiprocessing guard code
+# def main():
+#     obj = LeR()
+
+# if __name__ == '__main__':
+
 class LeR(LensGalaxyParameterDistribution):
     """Class to sample of lensed and unlensed events and calculate it's rates. Please note that parameters of the simulated events are stored in json file but not as an attribute of the class. This saves RAM memory. 
 
@@ -480,7 +488,7 @@ class LeR(LensGalaxyParameterDistribution):
         Function to print all the parameters.
         """
         # print all relevant functions and sampler priors
-        print("\n LeR set up params:")
+        print("\n # LeR set up params:")
         print(f'npool = {self.npool},')
         print(f'z_min = {self.z_min},')
         print(f'z_max = {self.z_max},')
@@ -493,23 +501,23 @@ class LeR(LensGalaxyParameterDistribution):
         if self.pdet:
             print(f'pdet_finder = {self.pdet},')
         print(f'json_file_names = {self.json_file_names},')
-        print(f'interpolator_directory = {self.interpolator_directory},')
-        print(f'ler_directory = {self.ler_directory},')
+        print(f"interpolator_directory = '{self.interpolator_directory}',")
+        print(f"ler_directory = '{self.ler_directory}',")
 
-        print("\n LeR also takes CBCSourceParameterDistribution class params as kwargs, as follows:")
+        print("\n # LeR also takes CBCSourceParameterDistribution class params as kwargs, as follows:")
         print(f"source_priors = {self.gw_param_sampler_dict['source_priors']},")
         print(f"source_priors_params = {self.gw_param_sampler_dict['source_priors_params']},")
         print(f"spin_zero = {self.gw_param_sampler_dict['spin_zero']},")
         print(f"spin_precession = {self.gw_param_sampler_dict['spin_precession']},")
         print(f"create_new_interpolator = {self.gw_param_sampler_dict['create_new_interpolator']},")
         
-        print("\n LeR also takes LensGalaxyParameterDistribution class params as kwargs, as follows:")
+        print("\n # LeR also takes LensGalaxyParameterDistribution class params as kwargs, as follows:")
         print(f"lens_type = '{self.gw_param_sampler_dict['lens_type']}',")
         print(f"lens_functions = {self.gw_param_sampler_dict['lens_functions']},")
         print(f"lens_priors = {self.gw_param_sampler_dict['lens_priors']},")
         print(f"lens_priors_params = {self.gw_param_sampler_dict['lens_priors_params']},")
         
-        print("\n LeR also takes ImageProperties class params as kwargs, as follows:")
+        print("\n # LeR also takes ImageProperties class params as kwargs, as follows:")
         print(f"n_min_images = {self.n_min_images},")
         print(f"n_max_images = {self.n_max_images},")
         print(f"geocent_time_min = {self.geocent_time_min},")
@@ -517,7 +525,7 @@ class LeR(LensGalaxyParameterDistribution):
         print(f"lens_model_list = {self.lens_model_list},")
         
         if self.gwsnr:
-            print("\n LeR also takes gwsnr.GWSNR params as kwargs, as follows:")
+            print("\n # LeR also takes gwsnr.GWSNR params as kwargs, as follows:")
             print(f"mtot_min = {self.snr_calculator_dict['mtot_min']},")
             print(f"mtot_max = {self.snr_calculator_dict['mtot_max']},")
             print(f"ratio_min = {self.snr_calculator_dict['ratio_min']},")
@@ -531,7 +539,6 @@ class LeR(LensGalaxyParameterDistribution):
             print(f"psds = {self.snr_calculator_dict['psds']},")
             print(f"ifos = {self.snr_calculator_dict['ifos']},")
             print(f"interpolator_dir = '{self.snr_calculator_dict['interpolator_dir']}',")
-            print(f"create_new_interpolator = {self.snr_calculator_dict['create_new_interpolator']},")
             print(f"gwsnr_verbose = {self.snr_calculator_dict['gwsnr_verbose']},")
             print(f"multiprocessing_verbose = {self.snr_calculator_dict['multiprocessing_verbose']},")
             print(f"mtot_cut = {self.snr_calculator_dict['mtot_cut']},")
@@ -712,6 +719,7 @@ class LeR(LensGalaxyParameterDistribution):
         # initialization of LensGalaxyParameterDistribution class
         # it also initializes the CBCSourceParameterDistribution and ImageProperties classes
         input_params = dict(
+            # LensGalaxyParameterDistribution class params
             z_min=self.z_min,
             z_max=self.z_max,
             cosmology=self.cosmo,
@@ -720,8 +728,13 @@ class LeR(LensGalaxyParameterDistribution):
             lens_functions= None,
             lens_priors=None,
             lens_priors_params=None,
+            # ImageProperties class params
+            n_min_images=2, 
+            n_max_images=4,
             geocent_time_min=1126259462.4,
             geocent_time_max=1126259462.4+365*24*3600*20,
+            lens_model_list=['EPL_NUMBA', 'SHEAR'],
+            # CBCSourceParameterDistribution class params
             source_priors=None,
             source_priors_params=None,
             spin_zero=True,
@@ -745,8 +758,11 @@ class LeR(LensGalaxyParameterDistribution):
             lens_functions=input_params["lens_functions"],
             lens_priors=input_params["lens_priors"],
             lens_priors_params=input_params["lens_priors_params"],
+            n_min_images=input_params["n_min_images"],
+            n_max_images=input_params["n_max_images"],
             geocent_time_min=input_params["geocent_time_min"],
             geocent_time_max=input_params["geocent_time_max"],
+            lens_model_list=input_params["lens_model_list"],
             source_priors=input_params["source_priors"],
             source_priors_params=input_params["source_priors_params"],
             spin_zero=input_params["spin_zero"],
@@ -792,6 +808,10 @@ class LeR(LensGalaxyParameterDistribution):
             gwsnr_verbose=False,
             multiprocessing_verbose=True,
             mtot_cut=True,
+            pdet=False,
+            snr_th=8.0,
+            snr_th_net=8.0,
+            ann_path_dict=None,
         )
         # if self.event_type == "BNS":
         #     input_params["mtot_max"]= 18.
@@ -819,6 +839,10 @@ class LeR(LensGalaxyParameterDistribution):
                     gwsnr_verbose=input_params["gwsnr_verbose"],
                     multiprocessing_verbose=input_params["multiprocessing_verbose"],
                     mtot_cut=input_params["mtot_cut"],
+                    pdet=input_params["pdet"],
+                    snr_th=input_params["snr_th"],
+                    snr_th_net=input_params["snr_th_net"],
+                    ann_path_dict=input_params["ann_path_dict"],
                 )
 
         self.snr = gwsnr.snr
@@ -909,36 +933,16 @@ class LeR(LensGalaxyParameterDistribution):
         output_path = os.path.join(self.ler_directory, output_jsonfile)
         print(f"unlensed params will be store in {output_path}")
 
-        # sampling in batches
-        if resume and os.path.exists(output_path):
-            # get sample from json file
-            self.dict_buffer = get_param_from_json(output_path)
-        else:
-            self.dict_buffer = None
-
-        batch_handler(
+        unlensed_param = batch_handler(
             size=size,
             batch_size=self.batch_size,
             sampling_routine=self.unlensed_sampling_routine,
             output_jsonfile=output_path,
             save_batch=save_batch,
             resume=resume,
+            param_name="unlensed parameters",
         )
 
-        if save_batch:
-            unlensed_param = get_param_from_json(output_path)
-        else:
-            # this if condition is required if there is nothing to save
-            if self.dict_buffer:
-                unlensed_param = self.dict_buffer.copy()
-                # store all params in json file
-                print(f"saving all unlensed_params in {output_path} ")
-                append_json(output_path, unlensed_param, replace=True)
-            else:
-                print("unlensed_params already sampled.")
-                unlensed_param = get_param_from_json(output_path)
-        self.dict_buffer = None  # save memory
-
         return unlensed_param
     
     def unlensed_sampling_routine(self, size, output_jsonfile, resume=False, save_batch=True):
@@ -978,17 +982,6 @@ class LeR(LensGalaxyParameterDistribution):
             pdet = self.pdet(gw_param_dict=unlensed_param)
             unlensed_param.update(pdet)
 
-        # adding batches
-        if not save_batch:
-            if self.dict_buffer is None:
-                self.dict_buffer = unlensed_param
-            else:
-                for key, value in unlensed_param.items():
-                    self.dict_buffer[key] = np.concatenate((self.dict_buffer[key], value))
-        else:
-            # store all params in json file
-            self.dict_buffer = append_json(file_name=output_jsonfile, new_dictionary=unlensed_param,  old_dictionary=self.dict_buffer, replace=not (resume))
-
         return unlensed_param
 
     def unlensed_rate(
@@ -1323,36 +1316,16 @@ class LeR(LensGalaxyParameterDistribution):
         output_path = os.path.join(self.ler_directory, output_jsonfile)
         print(f"lensed params will be store in {output_path}")
 
-        # sampling in batches
-        if resume and os.path.exists(output_path):
-            # get sample from json file
-            self.dict_buffer = get_param_from_json(output_path)
-        else:
-            self.dict_buffer = None
-
-        batch_handler(
+        lensed_param = batch_handler(
             size=size,
             batch_size=self.batch_size,
             sampling_routine=self.lensed_sampling_routine,
             output_jsonfile=output_path,
             save_batch=save_batch,
             resume=resume,
+            param_name="lensed parameters",
         )
 
-        if save_batch:
-            lensed_param = get_param_from_json(output_path)
-        else:
-            # this if condition is required if there is nothing to save
-            if self.dict_buffer:
-                lensed_param = self.dict_buffer.copy()
-                # store all params in json file
-                print(f"saving all lensed_params in {output_path} ")
-                append_json(output_path, lensed_param, replace=True)
-            else:
-                print("lensed_params already sampled.")
-                lensed_param = get_param_from_json(output_path)
-        self.dict_buffer = None  # save memory
-
         return lensed_param
     
     def lensed_sampling_routine(self, size, output_jsonfile, save_batch=True, resume=False):
@@ -1382,6 +1355,8 @@ class LeR(LensGalaxyParameterDistribution):
         print("sampling lensed params...")
         lensed_param = {}
 
+        # Some of the sample lensed events may not satisfy the strong lensing condition
+        # In that case, we will resample those events and replace the values with the corresponding indices
         while True:
             # get lensed params
             lensed_param_ = self.sample_lens_parameters(size=size)
@@ -1425,17 +1400,6 @@ class LeR(LensGalaxyParameterDistribution):
             )
             lensed_param.update(pdet)
 
-        # adding batches
-        if not save_batch:
-            if self.dict_buffer is None:
-                self.dict_buffer = lensed_param
-            else:
-                for key, value in lensed_param.items():
-                    self.dict_buffer[key] = np.concatenate((self.dict_buffer[key], value))
-        else:
-            # store all params in json file
-            self.dict_buffer = append_json(file_name=output_jsonfile, new_dictionary=lensed_param,  old_dictionary=self.dict_buffer, replace=not (resume))
-
         return lensed_param
 
     def lensed_rate(
@@ -1447,6 +1411,8 @@ class LeR(LensGalaxyParameterDistribution):
         output_jsonfile=None,
         nan_to_num=True,
         detectability_condition="step_function",
+        combine_image_snr=False,
+        snr_cut_for_combine_image_snr=4.0,
         snr_recalculation=False,
         snr_threshold_recalculation=[[4,4], [20,20]],
     ):
@@ -1519,7 +1485,8 @@ class LeR(LensGalaxyParameterDistribution):
         if snr_recalculation:
             lensed_param = self._recalculate_snr_lensed(lensed_param, snr_threshold_recalculation, num_img, total_events)
 
-        snr_hit = self._find_detectable_index_lensed(lensed_param, snr_threshold, pdet_threshold, num_img, detectability_condition)
+        # find index of detectable events
+        snr_hit = self._find_detectable_index_lensed(lensed_param, snr_threshold, pdet_threshold, num_img, detectability_condition, combine_image_snr=combine_image_snr, snr_cut_for_combine_image_snr=snr_cut_for_combine_image_snr)
 
         # montecarlo integration
         total_rate = self.rate_function(np.sum(snr_hit), total_events, param_type="lensed")
@@ -1619,7 +1586,7 @@ class LeR(LensGalaxyParameterDistribution):
 
         return lensed_param
 
-    def _find_detectable_index_lensed(self, lensed_param, snr_threshold, pdet_threshold, num_img, detectability_condition):
+    def _find_detectable_index_lensed(self, lensed_param, snr_threshold, pdet_threshold, num_img, detectability_condition, combine_image_snr=False, snr_cut_for_combine_image_snr=4.0):
         """
         Helper function to find the index of detectable events based on SNR or p_det.
 
@@ -1655,18 +1622,24 @@ class LeR(LensGalaxyParameterDistribution):
             snr_param = -np.sort(-snr_param, axis=1)  # sort snr in descending order
             snr_hit = np.full(len(snr_param), True)  # boolean array to store the result of the threshold condition
 
-            # for each row: choose a threshold and check if the number of images above threshold. Sum over the images. If sum is greater than num_img, then snr_hit = True 
-            # algorithm: 
-            # i) consider snr_threshold=[8,6] and num_img=[2,1] and first row of snr_param[0]=[12,8,6,1]. Note that the snr_param is sorted in descending order.
-            # ii) for loop runs wrt snr_threshold. idx_max = idx_max + num_img[i]
-            # iii) First iteration: snr_threshold=8 and num_img=2. In snr_param, column index 0 and 1 (i.e. 0:num_img[0]) are considered. The sum of snr_param[0, 0:2] > 8 is checked. If True, then snr_hit = True. 
-            # v) Second iteration: snr_threshold=6 and num_img=1. In snr_param, column index 2 (i.e. num_img[0]:num_img[1]) is considered. The sum of snr_param[0, 0:1] > 6 is checked. If True, then snr_hit = True.
-            j = 0
-            idx_max = 0
-            for i, snr_th in enumerate(snr_threshold):
-                idx_max = idx_max + num_img[i]
-                snr_hit = snr_hit & (np.sum((snr_param[:,j:idx_max] > snr_th), axis=1) >= num_img[i])
-                j = idx_max
+            if not combine_image_snr:
+                # for each row: choose a threshold and check if the number of images above threshold. Sum over the images. If sum is greater than num_img, then snr_hit = True 
+                # algorithm: 
+                # i) consider snr_threshold=[8,6] and num_img=[2,1] and first row of snr_param[0]=[12,8,6,1]. Note that the snr_param is sorted in descending order.
+                # ii) for loop runs wrt snr_threshold. idx_max = idx_max + num_img[i]
+                # iii) First iteration: snr_threshold=8 and num_img=2. In snr_param, column index 0 and 1 (i.e. 0:num_img[0]) are considered. The sum of snr_param[0, 0:2] > 8 is checked. If True, then snr_hit = True. 
+                # v) Second iteration: snr_threshold=6 and num_img=1. In snr_param, column index 2 (i.e. num_img[0]:num_img[1]) is considered. The sum of snr_param[0, 0:1] > 6 is checked. If True, then snr_hit = True.
+                j = 0
+                idx_max = 0
+                for i, snr_th in enumerate(snr_threshold):
+                    idx_max = idx_max + num_img[i]
+                    snr_hit = snr_hit & (np.sum((snr_param[:,j:idx_max] > snr_th), axis=1) >= num_img[i])
+                    j = idx_max
+            else:
+                # sqrt of the the sum of the squares of the snr of the images
+                snr_param[snr_param<snr_cut_for_combine_image_snr] = 0.0 # images with snr below snr_cut_for_combine_image_snr are not considered 
+                snr_param = np.sqrt(np.sum(snr_param[:,:np.sum(num_img)]**2, axis=1))
+                snr_hit = snr_param >= snr_threshold[0]
                 
         elif detectability_condition == "pdet":
             if "pdet_net" not in lensed_param:
@@ -1705,6 +1678,8 @@ class LeR(LensGalaxyParameterDistribution):
         lensed_param=None,
         snr_threshold_lensed=[8.0,8.0],
         num_img=[1,1],
+        combine_image_snr=False,
+        snr_cut_for_combine_image_snr=4.0,
         output_jsonfile_lensed=None,
         nan_to_num=True,
         detectability_condition="step_function",
@@ -1779,6 +1754,8 @@ class LeR(LensGalaxyParameterDistribution):
             output_jsonfile=output_jsonfile_lensed,
             nan_to_num=nan_to_num,
             detectability_condition=detectability_condition,
+            combine_image_snr=combine_image_snr,
+            snr_cut_for_combine_image_snr=snr_cut_for_combine_image_snr,
         )
         # calculate rate ratio
         rate_ratio = self.rate_ratio()
@@ -1893,7 +1870,7 @@ class LeR(LensGalaxyParameterDistribution):
         """
 
         # initial setup
-        n, events_total, output_path, meta_data_path, buffer_file = self._initial_setup_for_n_event_selection(meta_data_file, output_jsonfile, resume, batch_size)
+        n, events_total, output_path, meta_data_path, buffer_file, batch_size = self._initial_setup_for_n_event_selection(meta_data_file, output_jsonfile, resume, batch_size)
 
         # loop until n samples are collected
         while n < size:
@@ -1953,6 +1930,8 @@ class LeR(LensGalaxyParameterDistribution):
         snr_threshold=[8.0,8.0],
         pdet_threshold=0.5,
         num_img=[1,1],
+        combine_image_snr=False,
+        snr_cut_for_combine_image_snr=4.0,
         resume=False,
         detectability_condition="step_function",
         output_jsonfile="n_lensed_params_detectable.json",
@@ -2021,7 +2000,7 @@ class LeR(LensGalaxyParameterDistribution):
         """
 
         # initial setup
-        n, events_total, output_path, meta_data_path, buffer_file = self._initial_setup_for_n_event_selection(meta_data_file, output_jsonfile, resume, batch_size)
+        n, events_total, output_path, meta_data_path, buffer_file, batch_size = self._initial_setup_for_n_event_selection(meta_data_file, output_jsonfile, resume, batch_size)
 
         # re-analyse the provided snr_threshold and num_img
         snr_threshold, num_img = self._check_snr_threshold_lensed(snr_threshold, num_img)
@@ -2041,7 +2020,7 @@ class LeR(LensGalaxyParameterDistribution):
             if snr_recalculation:
                 lensed_param = self._recalculate_snr_lensed(lensed_param, snr_threshold_recalculation, num_img, total_events_in_this_iteration)
 
-            snr_hit = self._find_detectable_index_lensed(lensed_param, snr_threshold, pdet_threshold, num_img, detectability_condition)
+            snr_hit = self._find_detectable_index_lensed(lensed_param, snr_threshold, pdet_threshold, num_img, detectability_condition, combine_image_snr=combine_image_snr, snr_cut_for_combine_image_snr=snr_cut_for_combine_image_snr)
                     
             # store all params in json file
             self._save_detectable_params(output_jsonfile, lensed_param, snr_hit, key_file_name="n_lensed_detectable_events", nan_to_num=nan_to_num, verbose=False, replace_jsonfile=False)
@@ -2118,6 +2097,8 @@ class LeR(LensGalaxyParameterDistribution):
         if not resume:
             n = 0  # iterator
             events_total = 0
+            # the following file will be removed if it exists
+            print(f"removing {output_path} and {meta_data_path} if they exist")
             if os.path.exists(output_path):
                 os.remove(output_path)
             if os.path.exists(meta_data_path):
@@ -2136,7 +2117,7 @@ class LeR(LensGalaxyParameterDistribution):
         buffer_file = "params_buffer.json"
         print("collected number of detectable events = ", n)
 
-        return n, events_total, output_path, meta_data_path, buffer_file
+        return n, events_total, output_path, meta_data_path, buffer_file, batch_size
 
     def _trim_results_to_size(self, size, output_path, meta_data_path, param_type="unlensed"):
         """