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

ler/rates/ler.py

@@ -4,6 +4,7 @@ This module contains the main class for calculating the rates of detectable grav
 """
 
 import os
+# os.environ['OMP_NESTED'] = 'FALSE'
 import warnings
 warnings.filterwarnings("ignore")
 import logging
@@ -16,12 +17,6 @@ from ..lens_galaxy_population import LensGalaxyParameterDistribution
 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. 
 
@@ -59,7 +54,7 @@ class LeR(LensGalaxyParameterDistribution):
         def snr_finder(gw_param_dict):
             ...
             return optimal_snr_dict
-        where optimal_snr_dict.keys = ['optimal_snr_net']. Refer to `gwsnr` package's GWSNR.snr attribute for more details.
+        where optimal_snr_dict.keys = ['snr_net']. Refer to `gwsnr` package's GWSNR.snr attribute for more details.
     pdet_finder : `function`
         default pdet_finder = None.
         The rate calculation uses either the pdet_finder or the snr_finder to calculate the detectable events. The custom pdet finder function should follow the following signature:
@@ -165,7 +160,7 @@ class LeR(LensGalaxyParameterDistribution):
     |:meth:`~class_initialization`        | Function to initialize the       |
     |                                     | parent classes                   |
     +-------------------------------------+----------------------------------+
-    |:meth:`~gwsnr_intialization`         | Function to initialize the       |
+    |:meth:`~gwsnr_initialization`         | Function to initialize the       |
     |                                     | gwsnr class                      |
     +-------------------------------------+----------------------------------+
     |:meth:`~snr`                         | Function to get the snr with the |
@@ -460,6 +455,7 @@ class LeR(LensGalaxyParameterDistribution):
         **kwargs,
     ):
         
+        print("\nInitializing LeR class...\n")
         self.npool = npool
         self.z_min = z_min
         self.z_max = z_max
@@ -482,7 +478,7 @@ class LeR(LensGalaxyParameterDistribution):
             self.class_initialization(params=kwargs)
             # initialization self.snr and self.pdet from GWSNR class
             if not snr_finder and not pdet_finder:
-                self.gwsnr_intialization(params=kwargs)
+                self.gwsnr_initialization(params=kwargs)
                 self.gwsnr = True
                 self.pdet = pdet_finder
             else:
@@ -494,15 +490,15 @@ class LeR(LensGalaxyParameterDistribution):
             # store all the ler input parameters
             self.store_ler_params(output_jsonfile=self.json_file_names["ler_params"])
 
-        # if verbose, prevent anything from printing
+        # if not verbose, prevent anything from printing
         if verbose:
             initialization()
-            self.print_all_params()
+            self.print_all_params_ler()
         else:
             with contextlib.redirect_stdout(None):
                 initialization()
             
-    def print_all_params(self):
+    def print_all_params_ler(self):
         """
         Function to print all the parameters.
         """
@@ -533,14 +529,15 @@ class LeR(LensGalaxyParameterDistribution):
         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(f"lens_param_samplers = {self.gw_param_sampler_dict['lens_param_samplers']},")
+        print(f"lens_param_samplers_params = {self.gw_param_sampler_dict['lens_param_samplers_params']},")
         
         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},")
-        print(f"geocent_time_max = {self.geocent_time_max},")
+        print(f"time_window = {self.time_window},")
+        # print(f"geocent_time_min = {self.geocent_time_min},")
+        # print(f"geocent_time_max = {self.geocent_time_max},")
         print(f"lens_model_list = {self.lens_model_list},")
         
         if self.gwsnr:
@@ -604,14 +601,16 @@ class LeR(LensGalaxyParameterDistribution):
         print("axis_ratio_params = ", self.lens_param_samplers_params["axis_ratio"])
         print(f"axis_rotation_angle = '{self.lens_param_samplers['axis_rotation_angle']}'")
         print("axis_rotation_angle_params = ", self.lens_param_samplers_params["axis_rotation_angle"])
-        print(f"shear = '{self.lens_param_samplers['shear']}'")
-        print("shear_params = ", self.lens_param_samplers_params["shear"])
-        print(f"mass_density_spectral_index = '{self.lens_param_samplers['mass_density_spectral_index']}'")
-        print("mass_density_spectral_index_params = ", self.lens_param_samplers_params["mass_density_spectral_index"])
+        print(f"shear = '{self.lens_param_samplers['external_shear']}'")
+        print("shear_params = ", self.lens_param_samplers_params['external_shear'])
+        print(f"density_profile_slope = '{self.lens_param_samplers['density_profile_slope']}'")
+        print("density_profile_slope_params = ", self.lens_param_samplers_params["density_profile_slope"])
         # lens functions
         print("Lens functions:")
         print(f"strong_lensing_condition = '{self.lens_functions['strong_lensing_condition']}'")
         print(f"optical_depth = '{self.lens_functions['optical_depth']}'")
+        print(f"optical_depth_params = '{self.lens_functions_params['optical_depth']}'")
+        print(f"param_sampler_type = '{self.lens_functions['param_sampler_type']}'")
 
     @property
     def snr(self):
@@ -743,15 +742,18 @@ class LeR(LensGalaxyParameterDistribution):
             z_max=self.z_max,
             cosmology=self.cosmo,
             event_type=self.event_type,
-            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,
+            buffer_size=1000,
             # 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,
+            time_window=365*24*3600*20,
+            # 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,
@@ -775,13 +777,16 @@ class LeR(LensGalaxyParameterDistribution):
             event_type=input_params["event_type"],
             lens_type=input_params["lens_type"],
             lens_functions=input_params["lens_functions"],
-            lens_priors=input_params["lens_priors"],
-            lens_priors_params=input_params["lens_priors_params"],
+            lens_functions_params=input_params["lens_functions_params"],
+            lens_param_samplers=input_params["lens_param_samplers"],
+            lens_param_samplers_params=input_params["lens_param_samplers_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"],
+            time_window=input_params["time_window"],
+            # geocent_time_min=input_params["geocent_time_min"],
+            # geocent_time_max=input_params["geocent_time_max"],
             lens_model_list=input_params["lens_model_list"],
+            buffer_size=input_params["buffer_size"],
             source_priors=input_params["source_priors"],
             source_priors_params=input_params["source_priors_params"],
             spin_zero=input_params["spin_zero"],
@@ -792,11 +797,12 @@ class LeR(LensGalaxyParameterDistribution):
 
         self.gw_param_sampler_dict["source_priors"]=self.gw_param_samplers.copy()
         self.gw_param_sampler_dict["source_priors_params"]=self.gw_param_samplers_params.copy()
-        self.gw_param_sampler_dict["lens_priors"]=self.lens_param_samplers.copy()
-        self.gw_param_sampler_dict["lens_priors_params"]=self.lens_param_samplers_params.copy()
+        self.gw_param_sampler_dict["lens_param_samplers"]=self.lens_param_samplers.copy()
+        self.gw_param_sampler_dict["lens_param_samplers_params"]=self.lens_param_samplers_params.copy()
         self.gw_param_sampler_dict["lens_functions"]=self.lens_functions.copy()
+        self.gw_param_sampler_dict["lens_functions_params"]=self.lens_functions_params.copy()
 
-    def gwsnr_intialization(self, params=None):
+    def gwsnr_initialization(self, params=None):
         """
         Function to initialize the GWSNR class from the `gwsnr` package.
 
@@ -808,29 +814,55 @@ class LeR(LensGalaxyParameterDistribution):
         from gwsnr import GWSNR
 
         # initialization of GWSNR class
+        if 'mminbh' in self.gw_param_samplers_params['source_frame_masses']:
+            min_bh_mass = self.gw_param_samplers_params['source_frame_masses']['mminbh']
+        else:
+            min_bh_mass = 2.0
+
+        if 'mmaxbh' in self.gw_param_samplers_params['source_frame_masses']:
+            max_bh_mass = self.gw_param_samplers_params['source_frame_masses']['mmaxbh']
+        else:
+            max_bh_mass = 200.0
         input_params = dict(
+            # General settings
             npool=self.npool,
-            mtot_min=2.0,
-            mtot_max=200,
+            snr_method="interpolation_aligned_spins",
+            snr_type="optimal_snr",
+            gwsnr_verbose=True,
+            multiprocessing_verbose=True,
+            pdet_kwargs=None,
+            # Settings for interpolation grid
+            mtot_min=min_bh_mass*2,
+            mtot_max=max_bh_mass*2*(1+self.z_max) if max_bh_mass*2*(1+self.z_max)<500.0 else 500.0,
             ratio_min=0.1,
             ratio_max=1.0,
-            mtot_resolution=500,
-            ratio_resolution=50,
+            spin_max=0.99,
+            mtot_resolution=200,
+            ratio_resolution=20,
+            spin_resolution=10,
+            batch_size_interpolation=1000000,
+            interpolator_dir="./interpolator_pickle",
+            create_new_interpolator=False,
+            # GW signal settings
             sampling_frequency=2048.0,
             waveform_approximant="IMRPhenomD",
+            frequency_domain_source_model='lal_binary_black_hole',
             minimum_frequency=20.0,
-            snr_type="interpolation",
+            reference_frequency=None,
+            duration_max=None,
+            duration_min=None,
+            fixed_duration=None,
+            mtot_cut=False,
+            # Detector settings
             psds=None,
             ifos=None,
-            interpolator_dir=self.interpolator_directory,
-            create_new_interpolator=False,
-            gwsnr_verbose=False,
-            multiprocessing_verbose=True,
-            mtot_cut=True,
-            pdet=False,
-            snr_th=8.0,
-            snr_th_net=8.0,
+            noise_realization=None, # not implemented yet
+            # ANN settings
             ann_path_dict=None,
+            # Hybrid SNR recalculation settings
+            snr_recalculation=False,
+            snr_recalculation_range=[6,14],
+            snr_recalculation_waveform_approximant="IMRPhenomXPHM",
         )
         # if self.event_type == "BNS":
         #     input_params["mtot_max"]= 18.
@@ -851,40 +883,51 @@ class LeR(LensGalaxyParameterDistribution):
                 else:
                     raise ValueError("create_new_interpolator['gwsnr'] should be a boolean.")
             else:
-                raise ValueError("create_new_interpolator should be a boolean or a dictionary with 'gwsnr' key.")
+                input_params["create_new_interpolator"] = False
 
         # initialization of GWSNR class
         gwsnr = GWSNR(
-                    npool=input_params["npool"],
-                    mtot_min=input_params["mtot_min"],
-                    mtot_max=input_params["mtot_max"],
-                    ratio_min=input_params["ratio_min"],
-                    ratio_max=input_params["ratio_max"],
-                    mtot_resolution=input_params["mtot_resolution"],
-                    ratio_resolution=input_params["ratio_resolution"],
-                    sampling_frequency=input_params["sampling_frequency"],
-                    waveform_approximant=input_params["waveform_approximant"],
-                    minimum_frequency=input_params["minimum_frequency"],
-                    snr_type=input_params["snr_type"],
-                    psds=input_params["psds"],
-                    ifos=input_params["ifos"],
-                    interpolator_dir=input_params["interpolator_dir"],
-                    create_new_interpolator=input_params["create_new_interpolator"],
-                    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"],
-                )
+            npool=input_params["npool"],
+            snr_method=input_params["snr_method"],
+            snr_type=input_params["snr_type"],
+            gwsnr_verbose=input_params["gwsnr_verbose"],
+            multiprocessing_verbose=input_params["multiprocessing_verbose"],
+            pdet_kwargs=input_params["pdet_kwargs"],
+            mtot_min=input_params["mtot_min"],
+            mtot_max=input_params["mtot_max"],
+            ratio_min=input_params["ratio_min"],
+            ratio_max=input_params["ratio_max"],
+            spin_max=input_params["spin_max"],
+            mtot_resolution=input_params["mtot_resolution"],
+            ratio_resolution=input_params["ratio_resolution"],
+            spin_resolution=input_params["spin_resolution"],
+            batch_size_interpolation=input_params["batch_size_interpolation"],
+            interpolator_dir=input_params["interpolator_dir"],
+            create_new_interpolator=input_params["create_new_interpolator"],
+            sampling_frequency=input_params["sampling_frequency"],
+            waveform_approximant=input_params["waveform_approximant"],
+            frequency_domain_source_model=input_params["frequency_domain_source_model"],
+            minimum_frequency=input_params["minimum_frequency"],
+            reference_frequency=input_params["reference_frequency"],
+            duration_max=input_params["duration_max"],
+            duration_min=input_params["duration_min"],
+            fixed_duration=input_params["fixed_duration"],
+            mtot_cut=input_params["mtot_cut"],
+            psds=input_params["psds"],
+            ifos=input_params["ifos"],
+            noise_realization=input_params["noise_realization"],
+            ann_path_dict=input_params["ann_path_dict"],
+            snr_recalculation=input_params["snr_recalculation"],
+            snr_recalculation_range=input_params["snr_recalculation_range"],
+            snr_recalculation_waveform_approximant=input_params["snr_recalculation_waveform_approximant"],
+        )
 
-        self.snr = gwsnr.snr
+        self.snr = gwsnr.optimal_snr
         self.list_of_detectors = gwsnr.detector_list
-        self.snr_bilby = gwsnr.compute_bilby_snr
+        self.snr_bilby = gwsnr.optimal_snr_with_inner_product
         self.snr_calculator_dict["mtot_max"] = gwsnr.mtot_max
         self.snr_calculator_dict["psds"] = gwsnr.psds_list
-        #self.pdet = gwsnr.pdet
+        self.pdet = gwsnr.pdet
 
     def store_ler_params(self, output_jsonfile="ler_params.json"):
         """
@@ -922,12 +965,12 @@ class LeR(LensGalaxyParameterDistribution):
             for key, value in snr_calculator_dict.items():
                 snr_calculator_dict[key] = str(value)
             parameters_dict.update({"snr_calculator_dict": snr_calculator_dict})
-
-            file_name = output_jsonfile
-            append_json(self.ler_directory+"/"+file_name, parameters_dict, replace=True)
         except:
             # if snr_calculator is custom function
             pass
+        
+        file_name = output_jsonfile
+        append_json(self.ler_directory+"/"+file_name, parameters_dict, replace=True)
 
     def unlensed_cbc_statistics(
         self, size=None, resume=False, save_batch=False, output_jsonfile=None,
@@ -965,7 +1008,7 @@ class LeR(LensGalaxyParameterDistribution):
         output_jsonfile = output_jsonfile or self.json_file_names["unlensed_param"]
         self.json_file_names["unlensed_param"] = output_jsonfile
         output_path = os.path.join(self.ler_directory, output_jsonfile)
-        print(f"unlensed params will be store in {output_path}")
+        print(f"unlensed params will be stored in {output_path}")
 
         unlensed_param = batch_handler(
             size=size,
@@ -1006,10 +1049,11 @@ class LeR(LensGalaxyParameterDistribution):
         # get gw params
         print("sampling gw source params...")
         unlensed_param = self.sample_gw_parameters(size=size)
+
         # Get all of the signal to noise ratios
         if self.snr:
             print("calculating snrs...")
-            snrs = self.snr(gw_param_dict=unlensed_param)
+            snrs = self.snr(gw_param_dict=unlensed_param.copy())
             unlensed_param.update(snrs)
         elif self.pdet:
             print("calculating pdet...")
@@ -1146,7 +1190,7 @@ class LeR(LensGalaxyParameterDistribution):
             dictionary of unlensed GW source parameters.
         """
 
-        snr_param = unlensed_param["optimal_snr_net"]
+        snr_param = unlensed_param["snr_net"]
         idx_detectable = (snr_param > snr_threshold_recalculation[0]) & (snr_param < snr_threshold_recalculation[1])
         # reduce the size of the dict
         for key, value in unlensed_param.items():
@@ -1180,20 +1224,20 @@ class LeR(LensGalaxyParameterDistribution):
         """
 
         if self.snr:
-            if "optimal_snr_net" not in unlensed_param:
-                raise ValueError("'optimal_snr_net' not in unlensed param dict provided")
+            if "snr_net" not in unlensed_param:
+                raise ValueError("'snr_net' not in unlensed param dict provided")
             if detectability_condition == "step_function":
-                print("given detectability_condition == 'step_function'")
-                param = unlensed_param["optimal_snr_net"]
+                #print("given detectability_condition == 'step_function'")
+                param = unlensed_param["snr_net"]
                 threshold = snr_threshold
             elif detectability_condition == "pdet":
-                print("given detectability_condition == 'pdet'")
-                param = 1 - norm.cdf(snr_threshold - unlensed_param["optimal_snr_net"])
+                #print("given detectability_condition == 'pdet'")
+                param = 1 - norm.cdf(snr_threshold - unlensed_param["snr_net"])
                 unlensed_param["pdet_net"] = param
                 threshold = pdet_threshold
         elif self.pdet:
             if "pdet_net" in unlensed_param:
-                print("given detectability_condition == 'pdet'")
+                #print("given detectability_condition == 'pdet'")
                 param = unlensed_param["pdet_net"]
                 threshold = pdet_threshold
             else:
@@ -1409,11 +1453,12 @@ class LeR(LensGalaxyParameterDistribution):
             # check for invalid samples
             idx = lensed_param["n_images"] < 2
             
-            if np.sum(idx) == 0:
-                break
-            else:
-                print(f"Invalid sample found. Resampling {np.sum(idx)} lensed events...")
-                size = np.sum(idx)
+            # if np.sum(idx) == 0:
+            #     break
+            # else:
+            #     print(f"Invalid sample found. Resampling {np.sum(idx)} lensed events...")
+            #     size = np.sum(idx)
+            break
                 
         # Get all of the signal to noise ratios
         if self.snr:
@@ -1522,6 +1567,8 @@ class LeR(LensGalaxyParameterDistribution):
         # 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)
 
+        # select according to time delay
+
         # montecarlo integration
         total_rate = self.rate_function(np.sum(snr_hit), total_events, param_type="lensed")
 
@@ -1590,10 +1637,10 @@ class LeR(LensGalaxyParameterDistribution):
         snr_threshold_recalculation_max, _ = self._check_snr_threshold_lensed(snr_threshold_recalculation[1], num_img)
 
         # check optimal_snr_net is provided in dict
-        if "optimal_snr_net" not in lensed_param:
+        if "snr_net" not in lensed_param:
             raise ValueError("optimal_snr_net not provided in lensed_param dict. Exiting...")
 
-        snr_param = lensed_param["optimal_snr_net"]
+        snr_param = lensed_param["snr_net"]
         snr_param = -np.sort(-snr_param, axis=1)  # sort snr in descending order
             
         # 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
@@ -1648,11 +1695,11 @@ class LeR(LensGalaxyParameterDistribution):
             boolean array to store the result of the threshold condition.
         """
 
-        print(f"given detectability_condition == {detectability_condition}")
+        #print(f"given detectability_condition == {detectability_condition}")
         if detectability_condition == "step_function":
-            if "optimal_snr_net" not in lensed_param:
-                raise ValueError("'optimal_snr_net' not in lensed parm dict provided")
-            snr_param = lensed_param["optimal_snr_net"]
+            if "snr_net" not in lensed_param:
+                raise ValueError("'snr_net' not in lensed parm dict provided")
+            snr_param = lensed_param["snr_net"]
             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
 
@@ -1668,6 +1715,7 @@ class LeR(LensGalaxyParameterDistribution):
                 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])
+                    # select according to time delays
                     j = idx_max
             else:
                 # sqrt of the the sum of the squares of the snr of the images
@@ -1677,12 +1725,12 @@ class LeR(LensGalaxyParameterDistribution):
                 
         elif detectability_condition == "pdet":
             if "pdet_net" not in lensed_param:
-                if "optimal_snr_net" not in lensed_param:
-                    raise ValueError("'optimal_snr_net' or 'pdet_net' not in lensed parm dict provided")
+                if "snr_net" not in lensed_param:
+                    raise ValueError("'snr_net' or 'pdet_net' not in lensed parm dict provided")
                 else:
-                    print("calculating pdet using 'optimal_snr_net'...")
+                    print("calculating pdet using 'snr_net'...")
                     # pdet dimension is (size, n_max_images)
-                    snr_param = lensed_param["optimal_snr_net"]
+                    snr_param = lensed_param["snr_net"]
                     snr_param = -np.sort(-snr_param, axis=1)  # sort snr in descending order
 
                     # column index beyong np.sum(num_img)-1 are not considered
@@ -1841,13 +1889,17 @@ class LeR(LensGalaxyParameterDistribution):
         self,
         size=100,
         batch_size=None,
+        stopping_criteria=dict(
+            relative_diff_percentage=0.5,
+            number_of_last_batches_to_check=4,
+        ),
         snr_threshold=8.0,
         pdet_threshold=0.5,
         resume=False,
         output_jsonfile="n_unlensed_param_detectable.json",
         meta_data_file="meta_unlensed.json",
         detectability_condition="step_function",
-        trim_to_size=True,
+        trim_to_size=False,
         snr_recalculation=False,
         snr_threshold_recalculation=[4, 12],
     ):
@@ -1907,7 +1959,8 @@ class LeR(LensGalaxyParameterDistribution):
         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:
+        continue_condition = True
+        while continue_condition:
             # disable print statements
             with contextlib.redirect_stdout(None):
                 self.dict_buffer = None  # this is used to store the sampled unlensed_param in batches when running the sampling_routine
@@ -1932,7 +1985,31 @@ class LeR(LensGalaxyParameterDistribution):
             total_rate = self.rate_function(n, events_total, param_type="unlensed", verbose=False)
 
             # bookmark
-            self._append_meta_data(meta_data_path, n, events_total, total_rate)
+            buffer_dict = self._append_meta_data(meta_data_path, n, events_total, total_rate)
+
+            if isinstance(stopping_criteria, dict):
+                total_rates = np.array(buffer_dict['total_rate'])
+                limit = stopping_criteria['relative_diff_percentage']
+                num_a = stopping_criteria['number_of_last_batches_to_check']
+                if len(total_rates)>num_a:
+                    num_a = int(-1*(num_a))
+                    # num_b = int(num_a)
+                    percentage_diff = np.abs((total_rates[num_a:]-total_rates[-1])/total_rates[-1])*100
+                    print(f"percentage difference for the last {abs(num_a)} batches = {percentage_diff}")
+                    if np.any(percentage_diff>limit):
+                        continue_condition &= True
+                    else:
+                        print(rf"stopping criteria of rate relative difference of {limit}% reached. If you want to collect more events, reduce stopping_criteria['relative_diff_percentage'] or put stopping_criteria=None.")
+                        continue_condition &= False
+
+            if isinstance(size, int):
+                if n<size:
+                    continue_condition |= True
+                else:
+                    print(rf"Given size={size} reached")
+                    continue_condition |= False
+                    if stopping_criteria is None:
+                        continue_condition &= False
 
         print(f"stored detectable unlensed params in {output_path}")
         print(f"stored meta data in {meta_data_path}")
@@ -1960,6 +2037,10 @@ class LeR(LensGalaxyParameterDistribution):
     def selecting_n_lensed_detectable_events(
         self,
         size=100,
+        stopping_criteria=dict(
+            relative_diff_percentage=0.5,
+            number_of_last_batches_to_check=4,
+        ),
         batch_size=None,
         snr_threshold=[8.0,8.0],
         pdet_threshold=0.5,
@@ -1970,7 +2051,7 @@ class LeR(LensGalaxyParameterDistribution):
         detectability_condition="step_function",
         output_jsonfile="n_lensed_params_detectable.json",
         meta_data_file="meta_lensed.json",
-        trim_to_size=True,
+        trim_to_size=False,
         nan_to_num=False,
         snr_recalculation=False,
         snr_threshold_recalculation=[[4,4],[12,12]],
@@ -2039,7 +2120,8 @@ class LeR(LensGalaxyParameterDistribution):
         # re-analyse the provided snr_threshold and num_img
         snr_threshold, num_img = self._check_snr_threshold_lensed(snr_threshold, num_img)
 
-        while n < size:
+        continue_condition = True
+        while continue_condition:
             # disable print statements
             with contextlib.redirect_stdout(None):
                 self.dict_buffer = None  # this is used to store the sampled lensed_param in batches when running the sampling_routine
@@ -2064,7 +2146,30 @@ class LeR(LensGalaxyParameterDistribution):
             total_rate = self.rate_function(n, events_total, param_type="lensed", verbose=False)
 
             # save meta data
-            self._append_meta_data(meta_data_path, n, events_total, total_rate)
+            buffer_dict = self._append_meta_data(meta_data_path, n, events_total, total_rate)
+
+            if isinstance(stopping_criteria, dict):
+                total_rates = np.array(buffer_dict['total_rate'])
+                limit = stopping_criteria['relative_diff_percentage']
+                num_a = stopping_criteria['number_of_last_batches_to_check']
+
+                if len(total_rates)>num_a:
+                    num_a = int(-1*(num_a))
+                    # num_b = int(num_a)
+                    percentage_diff = np.abs((total_rates[num_a:]-total_rates[-1])/total_rates[-1])*100
+                    print(f"percentage difference for the last {abs(num_a)} batches = {percentage_diff}")
+                    if np.any(percentage_diff>limit):
+                        continue_condition &= True
+                    else:
+                        print(rf"stopping criteria of rate relative difference of {limit}% reached. If you want to collect more events, reduce stopping_criteria['relative_diff_percentage']")
+                        continue_condition &= False
+
+            if isinstance(size, int):
+                if n<size:
+                    continue_condition |= True
+                else:
+                    print(rf"Given size={size} reached")
+                    continue_condition |= False
 
         print(f"storing detectable lensed params in {output_path}")
         print(f"storing meta data in {meta_data_path}")
@@ -2084,7 +2189,7 @@ class LeR(LensGalaxyParameterDistribution):
             meta = get_param_from_json(meta_data_path)
             data["detectable_lensed_rate_per_year"] = meta["total_rate"][-1]
             data["detectability_condition_lensed"] = detectability_condition
-        append_json(self.ler_directory+"/"+self.json_file_names["ler_params"], data, replace=True)
+        buffer_dict = append_json(self.ler_directory+"/"+self.json_file_names["ler_params"], data, replace=True)
 
         return param_final
 
@@ -2232,12 +2337,14 @@ class LeR(LensGalaxyParameterDistribution):
 
         if os.path.exists(meta_data_path):
             try:
-                append_json(meta_data_path, meta_data, replace=False)
+                dict_ = append_json(meta_data_path, meta_data, replace=False)
             except:
-                append_json(meta_data_path, meta_data, replace=True)
+                dict_ = append_json(meta_data_path, meta_data, replace=True)
         else:
-            append_json(meta_data_path, meta_data, replace=True)
+            dict_ = append_json(meta_data_path, meta_data, replace=True)
 
         print("collected number of detectable events = ", n)
         print("total number of events = ", events_total)
-        print(f"total rate (yr^-1): {total_rate}")
+        print(f"total rate (yr^-1): {total_rate}")
+
+        return dict_

ler/utils/__init__.py

@@ -1,2 +1,4 @@
 from .utils import *
 from .plots import *
+from .gwsnr_training_data_generator import *
+from .function_interpolation import *