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

ler/utils/__init__.py

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

ler/utils/function_interpolation.py

@@ -0,0 +1,322 @@
+
+from .utils import interpolator_pickle_path, cubic_spline_interpolator, cubic_spline_interpolator2d_array, inverse_transform_sampler, pdf_cubic_spline_interpolator2d_array, save_pickle, load_pickle, inverse_transform_sampler2d
+import numpy as np
+from scipy.integrate import quad, cumulative_trapezoid
+from scipy.interpolate import  CubicSpline
+from scipy.stats import gaussian_kde
+from numba import njit
+
+
+class FunctionConditioning():
+
+    def __init__(self,
+        function=None,  # can also be an array of function values 
+        x_array=None,
+        conditioned_y_array=None,  # if this is not none, 2D interpolation will be used
+        y_array=None,
+        gaussian_kde=False,
+        gaussian_kde_kwargs={},
+        param_dict_given={},
+        directory='./interpolator_pickle',
+        sub_directory='default',
+        name='default',
+        create_new=False,
+        create_function=False,
+        create_function_inverse=False,
+        create_pdf=False,
+        create_rvs=False,
+        multiprocessing_function=False,
+        callback=None,
+    ):
+        create = self.create_decision_function(create_function, create_function_inverse, create_pdf, create_rvs)
+        self.info = param_dict_given
+        self.callback = callback
+        
+        if create:
+            # create_interpolator input list
+            input_list = [function, x_array, conditioned_y_array, create_function_inverse, create_pdf, create_rvs, multiprocessing_function]
+            input_list_kde = [x_array, y_array, gaussian_kde_kwargs]
+
+            # check first whether the directory, subdirectory and pickle exist
+            path_inv_cdf, it_exist = interpolator_pickle_path(
+                param_dict_given=param_dict_given,
+                directory=directory,
+                sub_directory=sub_directory,
+                interpolator_name=name,
+            )
+
+            # if the interpolator exists, load it
+            if create_new:
+                it_exist = False
+
+            if it_exist:
+                print(f"{name} interpolator will be loaded from {path_inv_cdf}")
+                # load the interpolator
+                interpolator = load_pickle(path_inv_cdf)
+            else:
+                print(f"{name} interpolator will be generated at {path_inv_cdf}")
+                if not gaussian_kde:
+                    interpolator = self.create_interpolator(*input_list)
+                    # save the interpolator
+                    save_pickle(path_inv_cdf, interpolator)
+                else:
+                    # gaussian KDE
+                    interpolator = self.create_gaussian_kde(*input_list_kde)
+                    save_pickle(path_inv_cdf, interpolator)
+
+            if not gaussian_kde:
+                x_array = interpolator['x_array']
+                z_array = interpolator['z_array']
+                conditioned_y_array = interpolator['conditioned_y_array']
+                y_array = None
+                function_spline = interpolator['function_spline']
+                function_inverse_spline = interpolator['function_inverse_spline']
+                pdf_norm_const = interpolator['pdf_norm_const']
+                cdf_values = interpolator['cdf_values']
+
+                if conditioned_y_array is None:
+                    # function is 1D
+                    self.function = njit(lambda x: cubic_spline_interpolator(x, function_spline, x_array)) if create_function else None
+                    # inverse function is 1D
+                    self.function_inverse = njit(lambda x: cubic_spline_interpolator(x, function_inverse_spline, z_array)) if create_function_inverse else None
+
+                    # pdf is 1D
+                    self.pdf = njit(lambda x: cubic_spline_interpolator(x, function_spline, x_array)/pdf_norm_const) if create_pdf else None
+                    # sampler is 1D
+                    self.rvs = njit(lambda size: inverse_transform_sampler(size, cdf_values, x_array)) if create_rvs else None
+                    
+                else:
+                    self.function = njit(lambda x, y: cubic_spline_interpolator2d_array(x, y, function_spline, x_array, conditioned_y_array)) if create_function else None
+
+                    self.function_inverse = njit(lambda x, y: cubic_spline_interpolator2d_array(x, y, function_inverse_spline, z_array, conditioned_y_array)) if create_function_inverse else None
+
+                    self.pdf = njit(lambda x, y: pdf_cubic_spline_interpolator2d_array(x, y, pdf_norm_const, function_spline, x_array, conditioned_y_array)) if create_pdf else None
+
+                    self.rvs = njit(lambda size, y: inverse_transform_sampler2d(size, y, cdf_values, x_array, conditioned_y_array)) if create_rvs else None
+
+                self.x_array = x_array
+                self.z_array = z_array
+                self.conditioned_y_array = conditioned_y_array
+                self.function_spline = function_spline
+                self.function_inverse_spline = function_inverse_spline
+                self.pdf_norm_const = pdf_norm_const
+                self.cdf_values = cdf_values
+            else:
+                x_array = interpolator['x_array']
+                y_array = interpolator['y_array']
+                kde_object = interpolator['kde_object']
+
+                
+                self.pdf = lambda x: kde_object.pdf(x) if create_pdf else None
+                if y_array is None:
+                    self.rvs = lambda size: kde_object.resample(size)[0] if create_rvs else None
+                else:
+                    self.rvs = lambda size: kde_object.resample(size) if create_rvs else None
+
+                self.x_array = x_array
+                self.y_array = y_array
+                self.kde_object = kde_object
+
+    def __call__(self, *args):
+        args = [np.array(arg).reshape(-1) if isinstance(arg, float) else arg for arg in args]
+        return getattr(self, self.callback)(*args)
+            
+
+    def create_decision_function(self, create_function, create_function_inverse, create_pdf, create_rvs):
+
+        decision_bool = True
+        if not isinstance(create_function, bool) and callable(create_function):
+            self.function = create_function
+            decision_bool = False
+        if not isinstance(create_function_inverse, bool) and callable(create_function_inverse):
+            self.function_inverse = create_function_inverse
+            decision_bool = False
+        if not isinstance(create_pdf, bool) and callable(create_pdf):
+            self.pdf = create_pdf
+            decision_bool = False
+        if not isinstance(create_rvs, bool) and callable(create_rvs):
+            self.rvs = create_rvs
+            decision_bool = False
+        
+        return decision_bool
+
+
+    def create_gaussian_kde(self, x_array, y_array, gaussian_kde_kwargs):
+
+        # 1d KDE
+        if y_array is None:
+            kde = gaussian_kde(x_array, **gaussian_kde_kwargs)
+        else:
+            data = np.vstack([x_array, y_array])
+            kde = gaussian_kde(data, **gaussian_kde_kwargs)
+
+        return {
+            'x_array': x_array,
+            'y_array': y_array,
+            'kde_object': kde,
+        }
+
+    def create_interpolator(self, function, x_array, conditioned_y_array, create_function_inverse, create_pdf, create_rvs, multiprocessing_function):
+
+        # function can be numpy array or callable
+        # x_array, z_array are 2D arrays if conditioned_y_array is not None
+        x_array, z_array, conditioned_y_array = self.create_z_array(x_array, function, conditioned_y_array, create_pdf, create_rvs, multiprocessing_function)
+        del function
+
+        function_spline = self.function_spline_generator(x_array, z_array, conditioned_y_array)
+
+        if create_function_inverse:
+            if conditioned_y_array is None:
+                idx_sort = np.argsort(z_array)
+            else:
+                idx_sort = np.argsort(z_array, axis=1)
+            x_array_ = x_array[idx_sort]
+            z_array_ = z_array[idx_sort]
+
+            # check z_array is strictly increasing
+            # if (not np.all(np.diff(z_array) > 0)) or (not np.all(np.diff(z_array) < 0)):
+            #     raise ValueError("z_array must be strictly increasing")
+            
+            function_inverse_spline = self.function_spline_generator(z_array_, x_array_, conditioned_y_array)
+        else:
+            function_inverse_spline = None
+
+        if create_pdf or create_rvs:
+            # cannot have -ve pdf
+            pdf_norm_const = self.pdf_norm_const_generator(x_array, function_spline, conditioned_y_array)
+
+            if create_rvs:
+                cdf_values = self.cdf_values_generator(x_array, z_array, conditioned_y_array)
+            else:
+                cdf_values = None
+        else:
+            pdf_norm_const = None
+            cdf_values = None
+            
+        return {
+            'x_array': x_array,
+            'z_array': z_array,
+            'conditioned_y_array': conditioned_y_array,
+            'function_spline': function_spline,
+            'function_inverse_spline': function_inverse_spline,
+            'pdf_norm_const': pdf_norm_const,
+            'cdf_values': cdf_values,
+        }
+
+    def create_z_array(self, x_array, function, conditioned_y_array, create_pdf, create_rvs, multiprocessing_function):
+
+        if callable(function):
+            # 1D
+            if conditioned_y_array is None:
+                z_array = function(x_array)
+                # remove nan values
+                idx = np.argwhere(np.isnan(z_array))
+                x_array = np.delete(x_array, idx)
+                z_array = np.delete(z_array, idx)
+            # 2D
+            else:
+                # check if x_array is 2D, if not, make it 2D of shape (len(conditioned_y_array), len(x_array))
+                if x_array.ndim == 1:
+                    x_array = np.array([x_array]*len(conditioned_y_array))
+
+                idx = np.argsort(conditioned_y_array)
+                conditioned_y_array = conditioned_y_array[idx]
+                # x_array is 2D here, each row corresponds to a different y value
+                x_array = x_array[idx]
+                # sort each row of x_array
+                x_array = np.sort(x_array, axis=1)
+
+                if multiprocessing_function:
+                    z_array = function(x_array, conditioned_y_array)
+                else:
+                    z_list = []
+                    for i, y in enumerate(conditioned_y_array):
+                        try:
+                            z_list.append(function(x_array[i], y*np.ones_like(x_array[i])))
+                        except:
+                            # print(x_array[i], y)
+                            z_list.append(function(x_array[i], y))
+                    z_array = np.array(z_list)
+        else:
+            if conditioned_y_array is None:
+                z_array = function
+                # remove nan values
+                idx = np.argwhere(np.isnan(z_array))
+                x_array = np.delete(x_array, idx)
+                z_array = np.delete(z_array, idx)
+            else:
+                if x_array.ndim == 1:
+                    x_array = np.array([x_array]*len(conditioned_y_array))
+                if function.ndim == 1:
+                    raise ValueError('function must be 2D array if conditioned_y_array is not None')
+                # row sort
+                idx = np.argsort(conditioned_y_array)
+                conditioned_y_array = conditioned_y_array[idx]
+                # x_array is 2D here, each row corresponds to a different y value
+                x_array = x_array[idx]
+                z_array = function[idx]
+
+                z_list = []
+                x_list = []
+                for i in range(len(conditioned_y_array)):
+                    # column sort
+                    idx = np.argsort(x_array[i])
+                    x_list.append(x_array[i][idx])
+                    z_list.append(z_array[i][idx])
+                x_array = np.array(x_list)
+                z_array = np.array(z_list)
+
+        # cannot have -ve pdf
+        if create_pdf or create_rvs:
+            z_array[z_array < 0.0] = 0.0
+
+        return x_array, z_array, conditioned_y_array
+
+    def cdf_values_generator(self, x_array, z_array, conditioned_y_array):
+        # 1D case
+        if conditioned_y_array is None:
+            # z_array[z_array<0.] = 0. # already done
+            cdf_values = cumulative_trapezoid(z_array, x_array, initial=0)
+            cdf_values = cdf_values/cdf_values[-1]
+        # 2D case
+        else:
+            cdf_values = []
+            for i, y in enumerate(conditioned_y_array):
+                z_array_ = z_array[i]
+                z_array_[z_array_<0.] = 0.
+                cdfs_ = cumulative_trapezoid(z_array_, x_array[i], initial=0)
+
+                cdf_values.append(cdfs_/cdfs_[-1])
+                # cdf_values.append(cdfs_)
+        
+        return np.array(cdf_values)
+    
+    def pdf_norm_const_generator(self, x_array, function_spline, conditioned_y_array):
+        # 1D case
+        if conditioned_y_array is None:
+            pdf_unorm = lambda x: cubic_spline_interpolator(np.array([x]), function_spline, x_array)
+
+            norm = quad(pdf_unorm, min(x_array), max(x_array))[0]
+            return norm
+        # 2D case
+        else:
+            norm = []
+            for i, y in enumerate(conditioned_y_array):
+                pdf_unorm = lambda x: cubic_spline_interpolator(np.array([x]), function_spline[i], x_array[i])
+
+                norm.append(quad(pdf_unorm, min(x_array[i]), max(x_array[i]))[0])
+
+            return np.array(norm)
+    
+    def function_spline_generator(self, x_array, z_array, conditioned_y_array):
+        # 1D case
+        if conditioned_y_array is None:
+            function_spline = CubicSpline(x_array, z_array).c
+        # 2D case
+        else:
+            function_spline = []
+            for i, y in enumerate(conditioned_y_array):
+                
+                function_spline.append(CubicSpline(x_array[i], z_array[i]).c)
+            
+        return np.array(function_spline)

ler/utils/gwsnr_training_data_generator.py

@@ -0,0 +1,233 @@
+import os
+import numpy as np
+import matplotlib.pyplot as plt
+import contextlib
+from gwsnr import GWSNR
+from .utils import append_json, get_param_from_json
+
+class TrainingDataGenerator():
+
+    def __init__(self,
+        npool=4,
+        z_min=0.0,
+        z_max=5.0,
+        verbose=True,
+        **kwargs,  # ler and gwsnr arguments
+    ):
+
+        self.npool = npool
+        self.z_min = z_min
+        self.z_max = z_max
+        self.verbose = verbose
+
+        self.ler_init_args = dict(
+            event_type="BBH",
+            cosmology=None,
+            ler_directory="./ler_data",
+            spin_zero=False,
+            spin_precession=True,
+            # gwsnr args
+            mtot_min=2*4.98, # 4.98 Mo is the minimum component mass of BBH systems in GWTC-3
+            mtot_max=2*112.5+10.0, # 112.5 Mo is the maximum component mass of BBH systems in GWTC-3. 10.0 Mo is added to avoid edge effects.
+            ratio_min=0.1,
+            ratio_max=1.0,
+            spin_max=0.99,
+            mtot_resolution=200,
+            ratio_resolution=20,
+            spin_resolution=10,
+            sampling_frequency=2048.0,
+            waveform_approximant="IMRPhenomXPHM",
+            minimum_frequency=20.0,
+            snr_type="interpolation_aligned_spins",
+            psds=None,
+            ifos=None,
+            interpolator_dir="./interpolator_pickle",
+            create_new_interpolator=False,
+            gwsnr_verbose=True,
+            multiprocessing_verbose=True,
+            mtot_cut=False,
+        )
+        self.ler_init_args.update(kwargs)
+
+    def gw_parameters_generator(self, 
+        size, 
+        batch_size=100000, 
+        snr_recalculation=True,
+        trim_to_size=False, 
+        verbose=True, 
+        replace=False, 
+        data_distribution_range = [0, 2, 4, 6, 8, 10, 12, 14, 16, 100],
+        output_jsonfile="gw_parameters.json",
+    ):
+
+        args = self.ler_init_args.copy()
+        if snr_recalculation:
+            snr_type = 'interpolation_aligned_spins'
+        else:
+            snr_type = 'inner_product'
+
+        # 
+        from ler.rates import GWRATES
+
+        # ler initialization
+        ler = GWRATES(
+            npool=self.npool,
+            z_min=self.z_min,
+            z_max=self.z_max,  # becareful with this value
+            verbose=self.verbose,
+            # ler
+            event_type=args['event_type'],
+            cosmology=args['cosmology'],
+            ler_directory=args['ler_directory'],
+            spin_zero=args['spin_zero'],
+            spin_precession=args['spin_precession'],
+            # gwsnr args
+            mtot_min=args['mtot_min'],
+            mtot_max=args['mtot_max'],
+            ratio_min=args['ratio_min'],
+            ratio_max=args['ratio_max'],
+            mtot_resolution=args['mtot_resolution'],
+            ratio_resolution=args['ratio_resolution'],
+            sampling_frequency=args['sampling_frequency'],
+            waveform_approximant=args['waveform_approximant'],
+            minimum_frequency=args['minimum_frequency'],
+            snr_type=snr_type,
+            psds=args['psds'],
+            ifos=args['ifos'],
+            interpolator_dir=args['interpolator_dir'],
+            create_new_interpolator=args['create_new_interpolator'],
+            gwsnr_verbose=args['gwsnr_verbose'],
+            multiprocessing_verbose=args['multiprocessing_verbose'],
+            mtot_cut=args['mtot_cut'],
+        )
+        ler.batch_size = batch_size
+
+        # path to save parameters
+        json_path = f"{args['ler_directory']}/{output_jsonfile}"
+        if replace:
+            if os.path.exists(json_path):
+                os.remove(json_path)
+            len_final = 0
+        else:
+            if os.path.exists(json_path):
+                gw_param = get_param_from_json(json_path)
+                len_final = len(gw_param['snr_net'])
+                print(f'current size of the json file: {len_final}\n')
+            else:
+                len_final = 0
+
+        print(f'total event to collect: {size}\n')
+        while len_final<size:
+            with contextlib.redirect_stdout(None):
+                gw_param = ler.gw_cbc_statistics(size=ler.batch_size, resume=False)
+
+            if data_distribution_range is not None:
+                gw_param = self.helper_data_distribution(gw_param, data_distribution_range)
+
+                if gw_param is None:
+                    continue
+
+                if snr_recalculation:
+                    snrs = ler.snr_bilby(gw_param_dict=gw_param)
+                    gw_param.update(snrs)
+
+                    gw_param = self.helper_data_distribution(gw_param, data_distribution_range)
+
+            if gw_param is None:
+                print("No data in one of the given range")
+                continue
+            # save the parameters
+            append_json(json_path, gw_param, replace=False);
+
+            # print(f"Collected number of events: {len_}")
+            len_final += len(gw_param['snr_net'])
+            if verbose:
+                print(f"Collected number of events: {len_final}")
+
+        if trim_to_size:
+            gw_param = get_param_from_json(json_path)
+            for key, value in gw_param.items():
+                gw_param[key] = value[:size]
+            append_json(json_path, gw_param, replace=True);
+            len_final = len(gw_param['snr_net'])
+
+        print(f"final size: {len_final}\n")
+        print(f"json file saved at: {json_path}\n")
+
+    def helper_data_distribution(self, gw_param, data_distribution_range):
+    # optimal SNR 
+        snr = np.array(gw_param['snr_net'])
+
+        idx_arr = []
+        snr_range = np.array(data_distribution_range)
+        len_ = len(snr_range) 
+        len_arr = []  # size of len_arr is len_-1  
+        for j in range(len_-1):
+            idx_ = np.argwhere((snr>=snr_range[j]) & (snr<snr_range[j+1])).flatten()
+            idx_arr.append(idx_)
+            len_arr.append(len(idx_))
+
+        idx_arr = np.array(idx_arr, dtype=object)
+        len_ref = min(len_arr)
+        
+        if len_ref==0:
+            print("No data in one of the given range")
+            return None
+        else:
+            gw_param_final = {}
+            for j, len_ in enumerate(len_arr):  # loop over snr range
+                idx_buffer = np.random.choice(idx_arr[j], len_ref, replace=False)
+
+                for key, value in gw_param.items(): 
+                    
+                    try:
+                        buffer_ = value[idx_buffer]
+                    except IndexError:
+                        print(f"IndexError")
+                        print(f"key: {key}, len(value): {len(value)}, len(idx_buffer): {len(idx_buffer)}")
+                        print(f"rerun the code again with: replace=False")
+                        return None
+                    if j==0:
+                        gw_param_final[key] = buffer_
+                    else:
+                        gw_param_final[key] = np.concatenate([gw_param_final[key], buffer_])
+
+            return gw_param_final
+
+    def combine_dicts(self,
+        file_name_list=None,
+        path_list=None, 
+        detector='L1',
+        parameter_list=['mass_1', 'mass_2', 'luminosity_distance', 'theta_jn', 'psi', 'geocent_time', 'ra', 'dec', 'a_1', 'a_2', 'tilt_1', 'tilt_2'],
+        output_jsonfile="combined_data.json",
+    ):
+
+        parameter_list += [detector]
+        combined_dict = {}
+
+        if file_name_list is not None:
+            path_list = [f"{self.ler_init_args['ler_directory']}/{file_name}" for file_name in file_name_list]
+        elif path_list is None:
+            print("Please provide either file_name_list or path_list")
+            return None
+
+        for path in path_list:
+            data = get_param_from_json(path)
+            for key, value in data.items():
+                if key in parameter_list:
+                    if key in combined_dict:
+                        combined_dict[key] = np.concatenate([combined_dict[key], value])
+                    else:
+                        combined_dict[key] = value
+
+        # if 'snr_net' is not in the combined_dict, we can calculate it
+        combined_dict['snr_net'] = combined_dict[detector]
+
+        json_path = f"{self.ler_init_args['ler_directory']}/{output_jsonfile}"
+        print(f"json file saved at: {json_path}\n")
+        append_json(json_path, combined_dict, replace=True);
+
+    def delete_json_file(self, path_list):
+        for path in path_list:
+            if os.path.exists(path):
+                os.remove(path)

ler/utils/plots.py

@@ -174,7 +174,7 @@ def relative_mu_dt_lensed(
     # get magnifications, time_delays and snr
     mu = lensed_param["magnifications"]
     dt = lensed_param["time_delays"]
-    snr = lensed_param["optimal_snr_net"]
+    snr = lensed_param["snr_net"]
     image_type = lensed_param["image_type"]
 
     # pair images wrt to image_type