ler 0.3.8__py3-none-any.whl → 0.4.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of ler might be problematic. Click here for more details.

ler/rates/ler.py CHANGED
@@ -15,7 +15,7 @@ from ..utils import load_json, append_json, get_param_from_json, batch_handler,
15
15
 
16
16
 
17
17
  class LeR(LensGalaxyParameterDistribution):
18
- """Class to calculate both the rates of lensed and unlensed events. 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.
18
+ """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.
19
19
 
20
20
  Parameters
21
21
  ----------
@@ -29,25 +29,54 @@ class LeR(LensGalaxyParameterDistribution):
29
29
  z_max : `float`
30
30
  maximum redshift.
31
31
  default z_max = 10.
32
- for popI_II, popIII, primordial, BNS z_max = 10., 40., 40., 2. respectively.
32
+ for popI_II, popIII, primordial, BNS z_max = 10., 40., 40., 5. respectively.
33
+ event_type : `str`
34
+ type of event to generate.
35
+ default event_type = 'BBH'. Other options are 'BNS', 'NSBH'.
33
36
  size : `int`
34
37
  number of samples for sampling.
35
- default size = 100000.
38
+ default size = 100000. To get stable rates, size should be large (>=1e6).
36
39
  batch_size : `int`
37
40
  batch size for SNR calculation.
38
41
  default batch_size = 50000.
39
42
  reduce the batch size if you are getting memory error.
40
- recommended batch_size = 50000, if size = 1000000.
41
- snr_finder : `str`
43
+ recommended batch_size = 200000, if size = 1000000.
44
+ cosmology : `astropy.cosmology`
45
+ cosmology to use for the calculation.
46
+ default cosmology = LambdaCDM(H0=70, Om0=0.3, Ode0=0.7).
47
+ snr_finder : `str` or `function`
42
48
  default snr_finder = 'gwsnr'.
43
49
  if None, the SNR will be calculated using the gwsnr package.
44
- if 'custom', the SNR will be calculated using a custom function.
45
- The custom function should have input and output as given in GWSNR.snr method.
50
+ if custom snr finder function is provided, the SNR will be calculated using a custom function. The custom function should follow the following signature:
51
+ def snr_finder(gw_param_dict):
52
+ ...
53
+ return optimal_snr_dict
54
+ where optimal_snr_dict.keys = ['optimal_snr_net']. Refer to `gwsnr` package's GWSNR.snr attribute for more details.
55
+ pdet_finder : `function`
56
+ default pdet_finder = None.
57
+ 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:
58
+ def pdet_finder(gw_param_dict):
59
+ ...
60
+ return pdet_net_dict
61
+ where pdet_net_dict.keys = ['pdet_net']. For example uses, refer to [GRB pdet example](https://ler.readthedocs.io/en/latest/examples/rates/grb%20detection%20rate.html).
62
+ list_of_detectors : `list`
63
+ list of detectors.
64
+ default list_of_detectors = ['H1', 'L1', 'V1']. This is used for lensed SNR calculation wrt to the detectors. Provide 'None' if you only need net SNR/Pdet. Refer to ImageProperties.get_lensed_snrs for more details.
46
65
  json_file_names: `dict`
47
66
  names of the json files to strore the necessary parameters.
48
- default json_file_names = {'ler_params': 'LeR_params.json', 'unlensed_param': 'unlensed_param.json', 'unlensed_param_detectable': 'unlensed_param_detectable.json'}.\n
67
+ default json_file_names = {'ler_params': 'LeR_params.json', 'unlensed_param': 'unlensed_param.json', 'unlensed_param_detectable': 'unlensed_param_detectable.json'}.
68
+ interpolator_directory : `str`
69
+ directory to store the interpolators.
70
+ default interpolator_directory = './interpolator_pickle'. This is used for storing the various interpolators related to `ler` and `gwsnr` package.
71
+ ler_directory : `str`
72
+ directory to store the parameters.
73
+ default ler_directory = './ler_data'. This is used for storing the parameters of the simulated events.
74
+ verbose : `bool`
75
+ default verbose = True.
76
+ if True, the function will print all chosen parameters.
77
+ Choose False to prevent anything from printing.
49
78
  kwargs : `keyword arguments`
50
- Note : kwargs takes input for initializing the :class:`~ler.lens_galaxy_population.LensGalaxyParameterDistribution`, :meth:`~gwsnr_intialization`.
79
+ Note : kwargs takes input for initializing the :class:`~ler.lens_galaxy_population.LensGalaxyParameterDistribution`, :class:`~ler.gw_source_population.CBCSourceParameterDistribution`, :class:`~ler.gw_source_population.CBCSourceRedshiftDistribution` and :class:`~ler.image_properties.ImageProperties` classes. If snr_finder='gwsnr', then kwargs also takes input for initializing the :class:`~gwsnr.GWSNR` class. Please refer to the respective classes for more details.
51
80
 
52
81
  Examples
53
82
  ----------
@@ -55,6 +84,9 @@ class LeR(LensGalaxyParameterDistribution):
55
84
  >>> ler = LeR()
56
85
  >>> unlensed_params = ler.unlensed_cbc_statistics();
57
86
  >>> ler.unlensed_rate();
87
+ >>> lensed_params = ler.lensed_cbc_statistics();
88
+ >>> ler.lensed_rate();
89
+ >>> ler.rate_ratio();
58
90
 
59
91
  Instance Attributes
60
92
  ----------
@@ -62,6 +94,8 @@ class LeR(LensGalaxyParameterDistribution):
62
94
  +-------------------------------------+----------------------------------+
63
95
  | Atrributes | Type |
64
96
  +=====================================+==================================+
97
+ |:attr:`~npool` | `int` |
98
+ +-------------------------------------+----------------------------------+
65
99
  |:attr:`~z_min` | `float` |
66
100
  +-------------------------------------+----------------------------------+
67
101
  |:attr:`~z_max` | `float` |
@@ -76,7 +110,11 @@ class LeR(LensGalaxyParameterDistribution):
76
110
  +-------------------------------------+----------------------------------+
77
111
  |:attr:`~json_file_names` | `dict` |
78
112
  +-------------------------------------+----------------------------------+
79
- |:attr:`~directory` | `str` |
113
+ |:attr:`~interpolator_directory` | `str` |
114
+ +-------------------------------------+----------------------------------+
115
+ |:attr:`~ler_directory` | `str` |
116
+ +-------------------------------------+----------------------------------+
117
+ |:attr:`~gwsnr` | `bool` |
80
118
  +-------------------------------------+----------------------------------+
81
119
  |:attr:`~gw_param_sampler_dict` | `dict` |
82
120
  +-------------------------------------+----------------------------------+
@@ -108,62 +146,79 @@ class LeR(LensGalaxyParameterDistribution):
108
146
  |:meth:`~snr` | Function to get the snr with the |
109
147
  | | given parameters. |
110
148
  +-------------------------------------+----------------------------------+
149
+ |:meth:`~snr_bilby` | Function to get the snr with the |
150
+ | | given parameters using inner- |
151
+ | | product method. |
152
+ +-------------------------------------+----------------------------------+
153
+ |:meth:`~pdet` | Function to get the pdet with |
154
+ | | the given parameters. |
155
+ +-------------------------------------+----------------------------------+
111
156
  |:meth:`~store_ler_params` | Function to store the all the |
112
157
  | | necessary parameters. |
113
158
  +-------------------------------------+----------------------------------+
114
159
  |:meth:`~unlensed_cbc_statistics` | Function to generate unlensed |
115
- | | GW source parameters. |
160
+ | | GW source parameters in batches. |
116
161
  +-------------------------------------+----------------------------------+
117
162
  |:meth:`~unlensed_sampling_routine` | Function to generate unlensed |
118
- | | GW source parameters. |
163
+ | | GW source parameters. It stores |
164
+ | | the parameters of the generated |
165
+ | | events in a json file. |
119
166
  +-------------------------------------+----------------------------------+
120
167
  |:meth:`~unlensed_rate` | Function to calculate the |
121
- | | unlensed rate. |
122
- +-------------------------------------+----------------------------------+
123
- |:meth:`~selecting_n_unlensed_detectable_events` |
124
- +-------------------------------------+----------------------------------+
125
- | | Function to select n unlensed |
126
- | | detectable events. |
168
+ | | unlensed rate. It also stores |
169
+ | | the parameters of the detectable |
170
+ | | unlesed events in a json file. |
127
171
  +-------------------------------------+----------------------------------+
128
172
  |:meth:`~lensed_cbc_statistics` | Function to generate lensed |
129
173
  | | GW source parameters. |
130
174
  +-------------------------------------+----------------------------------+
131
175
  |:meth:`~lensed_sampling_routine` | Function to generate lensed |
132
- | | GW source parameters. |
176
+ | | GW source parameters. It stores |
177
+ | | the parameters of the generated |
178
+ | | events in a json file. |
133
179
  +-------------------------------------+----------------------------------+
134
180
  |:meth:`~lensed_rate` | Function to calculate the |
135
- | | lensed rate. |
181
+ | | lensed rate. It also stores the |
182
+ | | parameters of the detectable |
183
+ | | lensed events in a json file. |
136
184
  +-------------------------------------+----------------------------------+
137
185
  |:meth:`~rate_ratio` | Function to calculate the rate |
138
- | | ratio. |
186
+ | | ratio between lensed and |
187
+ | | unlensed events. |
139
188
  +-------------------------------------+----------------------------------+
140
189
  |:meth:`~rate_comparision_with_rate_calculation |
141
190
  +-------------------------------------+----------------------------------+
142
- | | Function to compare the rates |
143
- | | calculated using LeR between |
144
- | | unlensed and lensed events. |
191
+ | | Function to calculate rates for |
192
+ | | unleesed and lensed events and |
193
+ | | compare it with the rate. It also|
194
+ | | stores the parameters of the |
195
+ | | detectable events in a json file.|
145
196
  +-------------------------------------+----------------------------------+
146
- |:meth:`~param_plot` | Function to plot the |
147
- | | distribution of various |
148
- | | parameters. |
197
+ |:meth:`~selecting_n_unlensed_detectable_events` |
149
198
  +-------------------------------------+----------------------------------+
150
- |:meth:`~relative_mu_dt_lensed` | Function to calculate the |
151
- | | relative magnification and |
152
- | | relative time-delay of lensed |
153
- | | events. |
199
+ | | Function to select n unlensed |
200
+ | | detectable events. It stores the |
201
+ | | parameters of the detectable |
202
+ | | unlesed events in a json file. |
154
203
  +-------------------------------------+----------------------------------+
155
- |:meth:`~relative_mu_dt_unlensed` | Function to calculate the |
156
- | | relative magnification and |
157
- | | relative time-delay of unlensed |
158
- | | events. |
204
+ |:meth:`~selecting_n_lensed_detectable_events` |
159
205
  +-------------------------------------+----------------------------------+
160
- |:meth:`~ mu_vs_dt_plot` | Function to plot the |
161
- | | relative magnification vs |
162
- | | relative time-delay. |
206
+ | | Function to select n lensed |
207
+ | | detectable events. It stores the |
208
+ | | parameters of the detectable |
209
+ | | lensed events in a json file. |
163
210
  +-------------------------------------+----------------------------------+
211
+
212
+ Note: `LeR` class also inherits all the instances from the :class:`~ler.lens_galaxy_population.LensGalaxyParameterDistribution` class. Please refer to the :class:`~ler.lens_galaxy_population.LensGalaxyParameterDistribution` class for more details.
164
213
  """
165
214
 
166
215
  # Attributes
216
+
217
+ npool = None
218
+ """``int`` \n
219
+ Number of logical cores to use.
220
+ """
221
+
167
222
  z_min = None
168
223
  """``float`` \n
169
224
  Minimum redshift of the source population
@@ -197,7 +252,7 @@ class LeR(LensGalaxyParameterDistribution):
197
252
 
198
253
  json_file_names = None
199
254
  """``dict`` \n
200
- Names of the json files to strore the necessary parameters.
255
+ Names of the json files to store the necessary parameters.
201
256
  """
202
257
 
203
258
  interpolator_directory = None
@@ -210,14 +265,14 @@ class LeR(LensGalaxyParameterDistribution):
210
265
  Directory to store the parameters.
211
266
  """
212
267
 
213
- gw_param_sampler_dict = None
214
- """``dict`` \n
215
- Dictionary of parameters to initialize the ``CBCSourceParameterDistribution`` class.
268
+ gwsnr = None
269
+ """``bool`` \n
270
+ If True, the SNR will be calculated using the gwsnr package.
216
271
  """
217
272
 
218
- lens_param_sampler_dict = None
273
+ gw_param_sampler_dict = None
219
274
  """``dict`` \n
220
- Dictionary of parameters to initialize the ``LensGalaxyParameterDistribution`` class.
275
+ Dictionary of parameters to initialize the ``CBCSourceParameterDistribution`` class.
221
276
  """
222
277
 
223
278
  snr_calculator_dict = None
@@ -230,6 +285,136 @@ class LeR(LensGalaxyParameterDistribution):
230
285
  List of detectors.
231
286
  """
232
287
 
288
+ unlensed_param = None
289
+ """``dict`` \n
290
+ Dictionary of unlensed GW source parameters. The included parameters and their units are as follows (for default settings):\n
291
+ +--------------------+--------------+--------------------------------------+
292
+ | Parameter | Units | Description |
293
+ +====================+==============+======================================+
294
+ | zs | | redshift of the source |
295
+ +--------------------+--------------+--------------------------------------+
296
+ | geocent_time | s | GPS time of coalescence |
297
+ +--------------------+--------------+--------------------------------------+
298
+ | ra | rad | right ascension |
299
+ +--------------------+--------------+--------------------------------------+
300
+ | dec | rad | declination |
301
+ +--------------------+--------------+--------------------------------------+
302
+ | phase | rad | phase of GW at reference frequency |
303
+ +--------------------+--------------+--------------------------------------+
304
+ | psi | rad | polarization angle |
305
+ +--------------------+--------------+--------------------------------------+
306
+ | theta_jn | rad | inclination angle |
307
+ +--------------------+--------------+--------------------------------------+
308
+ | luminosity_distance| Mpc | luminosity distance |
309
+ +--------------------+--------------+--------------------------------------+
310
+ | mass_1_source | Msun | mass_1 of the compact binary |
311
+ | | | (source frame) |
312
+ +--------------------+--------------+--------------------------------------+
313
+ | mass_2_source | Msun | mass_2 of the compact binary |
314
+ | | | (source frame) |
315
+ +--------------------+--------------+--------------------------------------+
316
+ | mass_1 | Msun | mass_1 of the compact binary |
317
+ | | | (detector frame) |
318
+ +--------------------+--------------+--------------------------------------+
319
+ | mass_2 | Msun | mass_2 of the compact binary |
320
+ | | | (detector frame) |
321
+ +--------------------+--------------+--------------------------------------+
322
+ | L1 | | optimal snr of L1 |
323
+ +--------------------+--------------+--------------------------------------+
324
+ | H1 | | optimal snr of H1 |
325
+ +--------------------+--------------+--------------------------------------+
326
+ | V1 | | optimal snr of V1 |
327
+ +--------------------+--------------+--------------------------------------+
328
+ | optimal_snr_net | | optimal snr of the network |
329
+ +--------------------+--------------+--------------------------------------+
330
+ """
331
+
332
+ unlensed_param_detectable = None
333
+ """``dict`` \n
334
+ Dictionary of detectable unlensed GW source parameters. It includes the same parameters as the :attr:`~unlensed_param` attribute.
335
+ """
336
+
337
+ lensed_param = None
338
+ """``dict`` \n
339
+ Dictionary of lens parameters, images parameters and lensed GW source parameters. The included parameters and their units are as follows (for default settings):\n
340
+ +------------------------------+-----------+-------------------------------+
341
+ | Parameter | Units | Description |
342
+ +==============================+===========+===============================+
343
+ | zl | | redshift of the lens |
344
+ +------------------------------+-----------+-------------------------------+
345
+ | zs | | redshift of the source |
346
+ +------------------------------+-----------+-------------------------------+
347
+ | sigma |km s^-1 | velocity dispersion |
348
+ +------------------------------+-----------+-------------------------------+
349
+ | q | | axis ratio |
350
+ +------------------------------+-----------+-------------------------------+
351
+ | theta_E | arcsec | Einstein radius |
352
+ +------------------------------+-----------+-------------------------------+
353
+ | phi | rad | axis rotation angle |
354
+ +------------------------------+-----------+-------------------------------+
355
+ | e1 | | ellipticity component 1 |
356
+ +------------------------------+-----------+-------------------------------+
357
+ | e2 | | ellipticity component 2 |
358
+ +------------------------------+-----------+-------------------------------+
359
+ | gamma1 | | shear component 1 |
360
+ +------------------------------+-----------+-------------------------------+
361
+ | gamma2 | | shear component 2 |
362
+ +------------------------------+-----------+-------------------------------+
363
+ | gamma | | shear |
364
+ +------------------------------+-----------+-------------------------------+
365
+ | ra | rad | right ascension |
366
+ +------------------------------+-----------+-------------------------------+
367
+ | dec | rad | declination |
368
+ +------------------------------+-----------+-------------------------------+
369
+ | phase | rad | phase of GW at reference freq |
370
+ +------------------------------+-----------+-------------------------------+
371
+ | psi | rad | polarization angle |
372
+ +------------------------------+-----------+-------------------------------+
373
+ | theta_jn | rad | inclination angle |
374
+ +------------------------------+-----------+-------------------------------+
375
+ | mass_1_source | Msun | mass_1 of the compact binary |
376
+ | | | (source frame) |
377
+ +------------------------------+-----------+-------------------------------+
378
+ | mass_2_source | Msun | mass_2 of the compact binary |
379
+ | | | (source frame) |
380
+ +------------------------------+-----------+-------------------------------+
381
+ | mass_1 | Msun | mass_1 of the compact binary |
382
+ | | | (detector frame) |
383
+ +------------------------------+-----------+-------------------------------+
384
+ | mass_2 | Msun | mass_2 of the compact binary |
385
+ | | | (detector frame) |
386
+ +------------------------------+-----------+-------------------------------+
387
+ | x0_image_positions | | x0 image positions |
388
+ +------------------------------+-----------+-------------------------------+
389
+ | x1_image_positions | | x1 image positions |
390
+ +------------------------------+-----------+-------------------------------+
391
+ | magnifications | | magnifications |
392
+ +------------------------------+-----------+-------------------------------+
393
+ | time_delays | | time delays |
394
+ +------------------------------+-----------+-------------------------------+
395
+ | image_type | | image type |
396
+ +------------------------------+-----------+-------------------------------+
397
+ | n_images | | number of images |
398
+ +------------------------------+-----------+-------------------------------+
399
+ | effective_luminosity_distance| Mpc | effective luminosity distance |
400
+ +------------------------------+-----------+-------------------------------+
401
+ | effective_geocent_time | s | effective GPS time of coalesc |
402
+ +------------------------------+-----------+-------------------------------+
403
+ | L1 | | optimal snr of L1 |
404
+ +------------------------------+-----------+-------------------------------+
405
+ | H1 | | optimal snr of H1 |
406
+ +------------------------------+-----------+-------------------------------+
407
+ | V1 | | optimal snr of V1 |
408
+ +------------------------------+-----------+-------------------------------+
409
+ | optimal_snr_net | | optimal snr of the network |
410
+ +------------------------------+-----------+-------------------------------+
411
+ """
412
+
413
+ lensed_param_detectable = None
414
+ """``dict`` \n
415
+ Dictionary of detectable lensed GW source parameters.
416
+ """
417
+
233
418
  def __init__(
234
419
  self,
235
420
  npool=int(4),
@@ -294,64 +479,63 @@ class LeR(LensGalaxyParameterDistribution):
294
479
  """
295
480
  Function to print all the parameters.
296
481
  """
297
-
298
482
  # print all relevant functions and sampler priors
299
483
  print("\n LeR set up params:")
300
- print("npool = ", self.npool)
301
- print("z_min = ", self.z_min)
302
- print("z_max = ", self.z_max)
303
- print(f"event_type = '{self.event_type}'")
304
- print("size = ", self.size)
305
- print("batch_size = ", self.batch_size)
306
- print("cosmology = ", self.cosmo)
484
+ print(f'npool = {self.npool},')
485
+ print(f'z_min = {self.z_min},')
486
+ print(f'z_max = {self.z_max},')
487
+ print(f"event_type = '{self.event_type}',")
488
+ print(f'size = {self.size},')
489
+ print(f'batch_size = {self.batch_size},')
490
+ print(f'cosmology = {self.cosmo},')
307
491
  if self.snr:
308
- print("snr_finder = ", self.snr)
492
+ print(f'snr_finder = {self.snr},')
309
493
  if self.pdet:
310
- print("pdet_finder = ", self.pdet)
311
- print("json_file_names = ", self.json_file_names)
312
- print("interpolator_directory = ", self.interpolator_directory)
313
- print("ler_directory = ", self.ler_directory)
494
+ print(f'pdet_finder = {self.pdet},')
495
+ print(f'json_file_names = {self.json_file_names},')
496
+ print(f'interpolator_directory = {self.interpolator_directory},')
497
+ print(f'ler_directory = {self.ler_directory},')
498
+
499
+ print("\n LeR also takes CBCSourceParameterDistribution class params as kwargs, as follows:")
500
+ print(f"source_priors = {self.gw_param_sampler_dict['source_priors']},")
501
+ print(f"source_priors_params = {self.gw_param_sampler_dict['source_priors_params']},")
502
+ print(f"spin_zero = {self.gw_param_sampler_dict['spin_zero']},")
503
+ print(f"spin_precession = {self.gw_param_sampler_dict['spin_precession']},")
504
+ print(f"create_new_interpolator = {self.gw_param_sampler_dict['create_new_interpolator']},")
505
+
506
+ print("\n LeR also takes LensGalaxyParameterDistribution class params as kwargs, as follows:")
507
+ print(f"lens_type = '{self.gw_param_sampler_dict['lens_type']}',")
508
+ print(f"lens_functions = {self.gw_param_sampler_dict['lens_functions']},")
509
+ print(f"lens_priors = {self.gw_param_sampler_dict['lens_priors']},")
510
+ print(f"lens_priors_params = {self.gw_param_sampler_dict['lens_priors_params']},")
511
+
512
+ print("\n LeR also takes ImageProperties class params as kwargs, as follows:")
513
+ print(f"n_min_images = {self.n_min_images},")
514
+ print(f"n_max_images = {self.n_max_images},")
515
+ print(f"geocent_time_min = {self.geocent_time_min},")
516
+ print(f"geocent_time_max = {self.geocent_time_max},")
517
+ print(f"lens_model_list = {self.lens_model_list},")
314
518
 
315
- print("\n LeR also takes CBCSourceParameterDistribution params as kwargs, as follows:")
316
- print("source_priors=", self.gw_param_sampler_dict["source_priors"])
317
- print("source_priors_params=", self.gw_param_sampler_dict["source_priors_params"])
318
- print("spin_zero=", self.gw_param_sampler_dict["spin_zero"])
319
- print("spin_precession=", self.gw_param_sampler_dict["spin_precession"])
320
- print("create_new_interpolator=", self.gw_param_sampler_dict["create_new_interpolator"])
321
-
322
- print("\n LeR also takes LensGalaxyParameterDistribution params as kwargs, as follows:")
323
- print(f"lens_type = '{self.gw_param_sampler_dict['lens_type']}'")
324
- print("lens_functions = ", self.gw_param_sampler_dict["lens_functions"])
325
- print("lens_priors = ", self.gw_param_sampler_dict["lens_priors"])
326
- print("lens_priors_params = ", self.gw_param_sampler_dict["lens_priors_params"])
327
-
328
- print("\n Image properties:")
329
- print("n_min_images = ", self.n_min_images)
330
- print("n_max_images = ", self.n_max_images)
331
- print("geocent_time_min = ", self.geocent_time_min)
332
- print("geocent_time_max = ", self.geocent_time_max)
333
- print("lens_model_list = ", self.lens_model_list)
334
-
335
519
  if self.gwsnr:
336
520
  print("\n LeR also takes gwsnr.GWSNR params as kwargs, as follows:")
337
- print("mtot_min = ", self.snr_calculator_dict["mtot_min"])
338
- print("mtot_max = ", self.snr_calculator_dict["mtot_max"])
339
- print("ratio_min = ", self.snr_calculator_dict["ratio_min"])
340
- print("ratio_max = ", self.snr_calculator_dict["ratio_max"])
341
- print("mtot_resolution = ", self.snr_calculator_dict["mtot_resolution"])
342
- print("ratio_resolution = ", self.snr_calculator_dict["ratio_resolution"])
343
- print("sampling_frequency = ", self.snr_calculator_dict["sampling_frequency"])
344
- print(f"waveform_approximant = '{self.snr_calculator_dict['waveform_approximant']}'")
345
- print("minimum_frequency = ", self.snr_calculator_dict["minimum_frequency"])
346
- print(f"snr_type = '{self.snr_calculator_dict['snr_type']}'")
347
- print("psds = ", self.snr_calculator_dict["psds"])
348
- print("ifos = ", self.snr_calculator_dict["ifos"])
349
- print("interpolator_dir = ", self.snr_calculator_dict["interpolator_dir"])
350
- print("create_new_interpolator = ", self.snr_calculator_dict["create_new_interpolator"])
351
- print("gwsnr_verbose = ", self.snr_calculator_dict["gwsnr_verbose"])
352
- print("multiprocessing_verbose = ", self.snr_calculator_dict["multiprocessing_verbose"])
353
- print("mtot_cut = ", self.snr_calculator_dict["mtot_cut"])
354
- del self.gwsnr
521
+ print(f"mtot_min = {self.snr_calculator_dict['mtot_min']},")
522
+ print(f"mtot_max = {self.snr_calculator_dict['mtot_max']},")
523
+ print(f"ratio_min = {self.snr_calculator_dict['ratio_min']},")
524
+ print(f"ratio_max = {self.snr_calculator_dict['ratio_max']},")
525
+ print(f"mtot_resolution = {self.snr_calculator_dict['mtot_resolution']},")
526
+ print(f"ratio_resolution = {self.snr_calculator_dict['ratio_resolution']},")
527
+ print(f"sampling_frequency = {self.snr_calculator_dict['sampling_frequency']},")
528
+ print(f"waveform_approximant = '{self.snr_calculator_dict['waveform_approximant']}',")
529
+ print(f"minimum_frequency = {self.snr_calculator_dict['minimum_frequency']},")
530
+ print(f"snr_type = '{self.snr_calculator_dict['snr_type']}',")
531
+ print(f"psds = {self.snr_calculator_dict['psds']},")
532
+ print(f"ifos = {self.snr_calculator_dict['ifos']},")
533
+ print(f"interpolator_dir = '{self.snr_calculator_dict['interpolator_dir']}',")
534
+ print(f"create_new_interpolator = {self.snr_calculator_dict['create_new_interpolator']},")
535
+ print(f"gwsnr_verbose = {self.snr_calculator_dict['gwsnr_verbose']},")
536
+ print(f"multiprocessing_verbose = {self.snr_calculator_dict['multiprocessing_verbose']},")
537
+ print(f"mtot_cut = {self.snr_calculator_dict['mtot_cut']},")
538
+ # del self.gwsnr
355
539
 
356
540
  print("\n For reference, the chosen source parameters are listed below:")
357
541
  print(f"merger_rate_density = '{self.gw_param_samplers['merger_rate_density']}'")
@@ -370,12 +554,12 @@ class LeR(LensGalaxyParameterDistribution):
370
554
  print("psi_params = ", self.gw_param_samplers_params["psi"])
371
555
  print(f"theta_jn = '{self.gw_param_samplers['theta_jn']}'")
372
556
  print("theta_jn_params = ", self.gw_param_samplers_params["theta_jn"])
373
- if self.spin_zero==False:
557
+ if self.spin_zero is False:
374
558
  print(f"a_1 = '{self.gw_param_samplers['a_1']}'")
375
559
  print("a_1_params = ", self.gw_param_samplers_params["a_1"])
376
560
  print(f"a_2 = '{self.gw_param_samplers['a_2']}'")
377
561
  print("a_2_params = ", self.gw_param_samplers_params["a_2"])
378
- if self.spin_precession==True:
562
+ if self.spin_precession is True:
379
563
  print(f"tilt_1 = '{self.gw_param_samplers['tilt_1']}'")
380
564
  print("tilt_1_params = ", self.gw_param_samplers_params["tilt_1"])
381
565
  print(f"tilt_2 = '{self.gw_param_samplers['tilt_2']}'")
@@ -517,14 +701,7 @@ class LeR(LensGalaxyParameterDistribution):
517
701
 
518
702
  def class_initialization(self, params=None):
519
703
  """
520
- Function to initialize the parent classes. List of relevant initialized instances, \n
521
- 1. self.sample_source_redshift
522
- 2. self.sample_unlensed_parameters
523
- 3. self.normalization_pdf_z
524
- 4. self.sample_lens_parameters
525
- 5. self.normalization_pdf_z_lensed
526
- 6. self.image_properties
527
- 7. self.get_lensed_snrs
704
+ Function to initialize the parent classes.
528
705
 
529
706
  Parameters
530
707
  ----------
@@ -586,7 +763,7 @@ class LeR(LensGalaxyParameterDistribution):
586
763
 
587
764
  def gwsnr_intialization(self, params=None):
588
765
  """
589
- Function to initialize the gwsnr class
766
+ Function to initialize the GWSNR class from the `gwsnr` package.
590
767
 
591
768
  Parameters
592
769
  ----------
@@ -649,7 +826,7 @@ class LeR(LensGalaxyParameterDistribution):
649
826
  self.snr_bilby = gwsnr.compute_bilby_snr
650
827
  self.snr_calculator_dict["mtot_max"] = gwsnr.mtot_max
651
828
  self.snr_calculator_dict["psds"] = gwsnr.psds_list
652
- # self.pdet = gwsnr.pdet
829
+ #self.pdet = gwsnr.pdet
653
830
 
654
831
  def store_ler_params(self, output_jsonfile="ler_params.json"):
655
832
  """
@@ -698,7 +875,7 @@ class LeR(LensGalaxyParameterDistribution):
698
875
  self, size=None, resume=False, save_batch=False, output_jsonfile=None,
699
876
  ):
700
877
  """
701
- Function to generate unlensed GW source parameters. This function also stores the parameters in json file.
878
+ Function to generate unlensed GW source parameters. This function calls the unlensed_sampling_routine function to generate the parameters in batches. The generated parameters are stored in a json file; and if save_batch=True, it keeps updating the file in batches.
702
879
 
703
880
  Parameters
704
881
  ----------
@@ -712,13 +889,12 @@ class LeR(LensGalaxyParameterDistribution):
712
889
  if True, the function will save the parameters in batches. if False, the function will save all the parameters at the end of sampling. save_batch=False is faster.
713
890
  output_jsonfile : `str`
714
891
  json file name for storing the parameters.
715
- default output_jsonfile = 'unlensed_params.json'.
892
+ default output_jsonfile = 'unlensed_params.json'. Note that this file will be stored in the self.ler_directory.
716
893
 
717
894
  Returns
718
895
  ----------
719
896
  unlensed_param : `dict`
720
- dictionary of unlensed GW source parameters.
721
- unlensed_param.keys() = ['zs', 'geocent_time', 'ra', 'dec', 'phase', 'psi', 'theta_jn', 'luminosity_distance', 'mass_1_source', 'mass_2_source', 'mass_1', 'mass_2', 'optimal_snr_net', 'L1', 'H1', 'V1']
897
+ dictionary of unlensed GW source parameters. Refer to :attr:`~unlensed_param` for details.
722
898
 
723
899
  Examples
724
900
  ----------
@@ -767,25 +943,26 @@ class LeR(LensGalaxyParameterDistribution):
767
943
 
768
944
  def unlensed_sampling_routine(self, size, output_jsonfile, resume=False, save_batch=True):
769
945
  """
770
- Function to generate unlensed GW source parameters. This function also stores the parameters in json file.
946
+ Function to generate unlensed GW source parameters. This function also stores the parameters in json file in the current batch if save_batch=True.
771
947
 
772
948
  Parameters
773
949
  ----------
774
950
  size : `int`
775
951
  number of samples.
776
952
  default size = 100000.
777
- resume : `bool`
778
- resume = False (default) or True.
779
- if True, the function will resume from the last batch.
780
953
  output_jsonfile : `str`
781
954
  json file name for storing the parameters.
782
- default output_jsonfile = 'unlensed_params.json'.
955
+ default output_jsonfile = 'unlensed_params.json'. Note that this file will be stored in the self.ler_directory.
956
+ resume : `bool`
957
+ resume = False (default) or True.
958
+ if True, it appends the new samples to the existing json file.
959
+ save_batch : `bool`
960
+ if True, the function will save the parameters in batches. if False, the function will save all the parameters at the end of sampling. save_batch=False is faster.
783
961
 
784
962
  Returns
785
963
  ----------
786
964
  unlensed_param : `dict`
787
- dictionary of unlensed GW source parameters.
788
- unlensed_param.keys() = ['zs', 'geocent_time', 'ra', 'dec', 'phase', 'psi', 'theta_jn', 'luminosity_distance', 'mass_1_source', 'mass_2_source', 'mass_1', 'mass_2', 'optimal_snr_net', 'L1', 'H1', 'V1']
965
+ dictionary of unlensed GW source parameters. Refer to :attr:`~unlensed_param` for details.
789
966
  """
790
967
 
791
968
  # get gw params
@@ -818,13 +995,19 @@ class LeR(LensGalaxyParameterDistribution):
818
995
  self,
819
996
  unlensed_param=None,
820
997
  snr_threshold=8.0,
998
+ pdet_threshold=0.5,
821
999
  output_jsonfile=None,
822
1000
  detectability_condition="step_function",
823
1001
  snr_recalculation=False,
824
- snr_threshold_recalculation=5.5,
1002
+ snr_threshold_recalculation=[4, 20],
825
1003
  ):
826
1004
  """
827
- Function to calculate the unlensed rate. This function also stores the parameters of the detectable events in json file.
1005
+ Function to calculate the unlensed rate. This function also stores the parameters of the detectable events in json file. There are two conditions for detectability: 'step_function' and 'pdet'.
1006
+
1007
+ 1. 'step_function': If two images have SNR>8.0, then the event is detectable. This is a step function. This is with the assumption that SNR function is provided and not None.
1008
+ 2. 'pdet':
1009
+ i) If self.pdet is None and self.snr is not None, then it will calculate the pdet from the snr. There is no hard cut for this pdet and can have value ranging from 0 to 1 near the threshold.
1010
+ ii) If self.pdet is not None, then it will use the generated pdet.
828
1011
 
829
1012
  Parameters
830
1013
  ----------
@@ -834,6 +1017,9 @@ class LeR(LensGalaxyParameterDistribution):
834
1017
  snr_threshold : `float`
835
1018
  threshold for detection signal to noise ratio.
836
1019
  e.g. snr_threshold = 8.
1020
+ pdet_threshold : `float`
1021
+ threshold for detection probability.
1022
+ e.g. pdet_threshold = 0.5.
837
1023
  output_jsonfile : `str`
838
1024
  json file name for storing the parameters of the detectable events.
839
1025
  default output_jsonfile = 'unlensed_params_detectable.json'.
@@ -842,18 +1028,18 @@ class LeR(LensGalaxyParameterDistribution):
842
1028
  default detectability_condition = 'step_function'.
843
1029
  other options are 'pdet'.
844
1030
  snr_recalculation : `bool`
845
- if True, the SNR of centain events (snr>snr_threshold_recalculation)will be recalculate with 'inner product'. This is useful when the snr is calculated with 'ann' method.
1031
+ if True, the SNR of centain events (snr>snr_threshold_recalculation)will be recalculate with 'inner-product' method. This is useful when the snr is calculated with 'ann' method of `gwsnr`.
846
1032
  default snr_recalculation = False.
847
- snr_threshold_recalculation : `float`
848
- threshold for recalculation of detection signal to noise ratio.
1033
+ snr_threshold_recalculation : `list`
1034
+ lower and upper threshold for recalculation of detection signal to noise ratio.
1035
+ default snr_threshold_recalculation = [4, 20].
849
1036
 
850
1037
  Returns
851
1038
  ----------
852
1039
  total_rate : `float`
853
1040
  total unlensed rate (Mpc^-3 yr^-1).
854
1041
  unlensed_param : `dict`
855
- dictionary of unlensed GW source parameters of the detectable events.
856
- unlensed_param.keys() = ['zs', 'geocent_time', 'ra', 'dec', 'phase', 'psi', 'theta_jn', 'luminosity_distance', 'mass_1_source', 'mass_2_source', 'mass_1', 'mass_2', 'optimal_snr_net', 'L1', 'H1', 'V1']
1042
+ dictionary of unlensed GW source parameters of the detectable events. Refer to :attr:`~unlensed_param` for details.
857
1043
 
858
1044
  Examples
859
1045
  ----------
@@ -863,39 +1049,112 @@ class LeR(LensGalaxyParameterDistribution):
863
1049
  >>> total_rate, unlensed_param_detectable = ler.unlensed_rate()
864
1050
  """
865
1051
 
866
- # call self.json_file_names["ler_param"] and for adding the final results
867
- data = load_json(self.ler_directory+"/"+self.json_file_names["ler_params"])
1052
+ unlensed_param = self._load_param(unlensed_param, param_type="unlensed")
1053
+ total_events = len(unlensed_param["zs"])
1054
+
1055
+ # below is use when the snr is calculated with 'ann' method of `gwsnr`
1056
+ if snr_recalculation:
1057
+ unlensed_param = self._recalculate_snr_unlensed(unlensed_param, snr_threshold_recalculation)
1058
+
1059
+ # find index of detectable events
1060
+ idx_detectable = self._find_detectable_index_unlensed(unlensed_param, snr_threshold, pdet_threshold, detectability_condition)
1061
+
1062
+ detectable_events = np.sum(idx_detectable)
1063
+ # montecarlo integration
1064
+ # The total rate R = norm <Theta(rho-rhoc)>
1065
+ total_rate = self.rate_function(detectable_events, total_events, param_type="unlensed")
1066
+
1067
+ # store all detectable params in json file
1068
+ self._save_detectable_params(output_jsonfile, unlensed_param, idx_detectable, key_file_name="unlensed_param_detectable", nan_to_num=False, verbose=True, replace_jsonfile=True)
1069
+
1070
+ # append ler_param and save it
1071
+ self._append_ler_param(total_rate, detectability_condition)
1072
+
1073
+ return total_rate, unlensed_param
1074
+
1075
+ def _load_param(self, param, param_type="unlensed"):
1076
+ """
1077
+ Helper function to load or copy unlensed parameters.
868
1078
 
869
- # get gw params from json file if not provided
870
- if unlensed_param is None:
871
- unlensed_param = self.json_file_names["unlensed_param"]
872
- if type(unlensed_param) == str:
873
- self.json_file_names["unlensed_param"] = unlensed_param
874
- path_ = self.ler_directory+"/"+unlensed_param
875
- print(f"getting unlensed_params from json file {path_}...")
876
- unlensed_param = get_param_from_json(self.ler_directory+"/"+unlensed_param)
1079
+ Parameters
1080
+ ----------
1081
+ param : `dict` or `str`
1082
+ dictionary of unlensed/lensed parameters or json file name.
1083
+ param_type : `str`
1084
+ type of parameters.
1085
+ default param_type = 'unlensed'. Other options is 'lensed'.
1086
+
1087
+ Returns
1088
+ ----------
1089
+ param : `dict`
1090
+ dictionary of unlensed/lensed parameters.
1091
+ """
1092
+
1093
+ param_type = param_type+"_param"
1094
+ if param is None:
1095
+ param = self.json_file_names[param_type]
1096
+ if isinstance(param, str):
1097
+ path_ = self.ler_directory + "/" + param
1098
+ print(f"Getting {param_type} from json file {path_}...")
1099
+ return get_param_from_json(path_)
877
1100
  else:
878
- print("using provided unlensed_param dict...")
879
- unlensed_param = unlensed_param.copy()
1101
+ print("Using provided {param_type} dict...")
1102
+ return param.copy()
880
1103
 
881
- # recalculate snr if required
882
- # this ensures that the snr is recalculated for the detectable events
883
- # with inner product
884
- total_events = len(unlensed_param["zs"])
885
- if snr_recalculation:
886
- # select only above centain snr threshold
887
- param = unlensed_param["optimal_snr_net"]
888
- idx_detectable = param > snr_threshold_recalculation
889
- # reduce the size of the dict
890
- for key, value in unlensed_param.items():
891
- unlensed_param[key] = value[idx_detectable]
892
- # recalculate more accurate snrs
893
- snrs = self.snr_bilby(gw_param_dict=unlensed_param)
894
- unlensed_param.update(snrs)
1104
+ def _recalculate_snr_unlensed(self, unlensed_param, snr_threshold_recalculation):
1105
+ """
1106
+ Recalculates SNR for events where the initial SNR is above a given threshold.
1107
+
1108
+ Parameters
1109
+ ----------
1110
+ unlensed_param : `dict`
1111
+ dictionary of unlensed GW source parameters.
1112
+ snr_threshold_recalculation : `list`
1113
+ lower and upper threshold for recalculation of detection signal to noise ratio.
1114
+ default snr_threshold_recalculation = [4, 20].
1115
+
1116
+ Returns
1117
+ ----------
1118
+ unlensed_param : `dict`
1119
+ dictionary of unlensed GW source parameters.
1120
+ """
1121
+
1122
+ snr_param = unlensed_param["optimal_snr_net"]
1123
+ idx_detectable = (snr_param > snr_threshold_recalculation[0]) & (snr_param < snr_threshold_recalculation[1])
1124
+ # reduce the size of the dict
1125
+ for key, value in unlensed_param.items():
1126
+ unlensed_param[key] = value[idx_detectable]
1127
+ # recalculate more accurate snrs
1128
+ snrs = self.snr_bilby(gw_param_dict=unlensed_param)
1129
+ unlensed_param.update(snrs)
1130
+ return unlensed_param
1131
+
1132
+ def _find_detectable_index_unlensed(self, unlensed_param, snr_threshold, pdet_threshold, detectability_condition):
1133
+ """
1134
+ Find the index of detectable events based on SNR or p_det.
1135
+
1136
+ Parameters
1137
+ ----------
1138
+ unlensed_param : `dict`
1139
+ dictionary of unlensed GW source parameters.
1140
+ snr_threshold : `float`
1141
+ threshold for detection signal to noise ratio.
1142
+ pdet_threshold : `float`
1143
+ threshold for detection probability.
1144
+ detectability_condition : `str`
1145
+ detectability condition.
1146
+ default detectability_condition = 'step_function'.
1147
+ other options are 'pdet'.
1148
+
1149
+ Returns
1150
+ ----------
1151
+ idx_detectable : `numpy.ndarray`
1152
+ index of detectable events.
1153
+ """
895
1154
 
896
1155
  if self.snr:
897
1156
  if "optimal_snr_net" not in unlensed_param:
898
- raise ValueError("'optimal_snr_net' not in unlensed parm dict provided")
1157
+ raise ValueError("'optimal_snr_net' not in unlensed param dict provided")
899
1158
  if detectability_condition == "step_function":
900
1159
  print("given detectability_condition == 'step_function'")
901
1160
  param = unlensed_param["optimal_snr_net"]
@@ -904,55 +1163,141 @@ class LeR(LensGalaxyParameterDistribution):
904
1163
  print("given detectability_condition == 'pdet'")
905
1164
  param = 1 - norm.cdf(snr_threshold - unlensed_param["optimal_snr_net"])
906
1165
  unlensed_param["pdet_net"] = param
907
- threshold = 0.5
1166
+ threshold = pdet_threshold
908
1167
  elif self.pdet:
909
1168
  if "pdet_net" in unlensed_param:
910
1169
  print("given detectability_condition == 'pdet'")
911
1170
  param = unlensed_param["pdet_net"]
912
- threshold = 0.5
1171
+ threshold = pdet_threshold
913
1172
  else:
914
- raise ValueError("'pdet_net' not in unlensed parm dict provided")
1173
+ raise ValueError("'pdet_net' not in unlensed param dict provided")
915
1174
 
916
1175
  idx_detectable = param > threshold
917
- detectable_events = np.sum(idx_detectable)
918
- # montecarlo integration
919
- # The total rate R = norm <Theta(rho-rhoc)>
920
- total_rate = self.normalization_pdf_z * detectable_events/total_events
921
- print(f"total unlensed rate (yr^-1) (with step function): {total_rate}")
922
- print(f"number of simulated unlensed detectable events: {detectable_events}")
923
- print(f"number of all simulated unlensed events: {total_events}")
1176
+ return idx_detectable
1177
+
1178
+ def rate_function(self, detectable_size, total_size, param_type="unlensed", verbose=True):
1179
+ """
1180
+ General helper function to calculate the rate for unlensed and lensed events.
1181
+
1182
+ Parameters
1183
+ ----------
1184
+ detectable_size : `int`
1185
+ number of detectable events.
1186
+ total_size : `int`
1187
+ total number of events.
1188
+ param_type : `str`
1189
+ type of parameters.
1190
+
1191
+ Returns
1192
+ ----------
1193
+ rate : `float`
1194
+ rate of the events.
1195
+
1196
+ Examples
1197
+ ----------
1198
+ >>> from ler.rates import LeR
1199
+ >>> ler = LeR()
1200
+ >>> rate = ler.rate_function(detectable_size=100, total_size=1000)
1201
+ """
1202
+
1203
+ if param_type == "unlensed":
1204
+ normalization = self.normalization_pdf_z
1205
+ elif param_type == "lensed":
1206
+ normalization = self.normalization_pdf_z_lensed
1207
+ rate = normalization * detectable_size / total_size
1208
+
1209
+ if verbose:
1210
+ print(f"total {param_type} rate (yr^-1): {rate}")
1211
+ print(f"number of simulated {param_type} detectable events: {detectable_size}")
1212
+ print(f"number of simulated all {param_type} events: {total_size}")
1213
+
1214
+ return rate
1215
+
1216
+ def _save_detectable_params(self,
1217
+ output_jsonfile,
1218
+ param,
1219
+ idx_detectable,
1220
+ key_file_name="unlensed_param_detectable",
1221
+ nan_to_num=False,
1222
+ verbose=True,
1223
+ replace_jsonfile=True,
1224
+ ):
1225
+ """
1226
+ Helper function to save the detectable parameters in json file.
1227
+
1228
+ Parameters
1229
+ ----------
1230
+ output_jsonfile : `str`
1231
+ json file name for storing the parameters of the detectable events. This is stored in the self.ler_directory.
1232
+ param : `dict`
1233
+ dictionary of GW source parameters.
1234
+ idx_detectable : `numpy.ndarray`
1235
+ index of detectable events.
1236
+ key_file_name : `str`
1237
+ key name for the json file to be added in self.json_file_names.
1238
+ nan_to_num : `bool`
1239
+ if True, it will replace nan with 0.
1240
+ default nan_to_num = False.
1241
+ verbose : `bool`
1242
+ if True, it will print the path of the json file.
1243
+ default verbose = True.
1244
+ replace_jsonfile : `bool`
1245
+ if True, it will replace the json file. If False, it will append the json file.
1246
+ """
924
1247
 
925
1248
  # store all detectable params in json file
926
- for key, value in unlensed_param.items():
927
- unlensed_param[key] = np.array(value)[idx_detectable]
1249
+ if nan_to_num:
1250
+ for key, value in param.items():
1251
+ param[key] = np.nan_to_num(value[idx_detectable])
1252
+ else:
1253
+ for key, value in param.items():
1254
+ param[key] = value[idx_detectable]
928
1255
 
929
1256
  # store all detectable params in json file
930
1257
  if output_jsonfile is None:
931
- output_jsonfile = self.json_file_names["unlensed_param_detectable"]
1258
+ output_jsonfile = self.json_file_names[key_file_name]
932
1259
  else:
933
- self.json_file_names["unlensed_param_detectable"] = output_jsonfile
934
- path_ = self.ler_directory+"/"+output_jsonfile
935
- print(f"storing detectable unlensed params in {path_}")
936
- append_json(self.ler_directory+"/"+output_jsonfile, unlensed_param, replace=True)
1260
+ self.json_file_names[key_file_name] = output_jsonfile
1261
+
1262
+ output_path = self.ler_directory+"/"+output_jsonfile
1263
+ if verbose:
1264
+ print(f"storing detectable params in {output_path}")
1265
+ append_json(output_path, param, replace=replace_jsonfile)
1266
+
1267
+ def _append_ler_param(self, total_rate, detectability_condition, param_type="unlensed"):
1268
+ """
1269
+ Helper function to append the final results, total_rate, in the json file.
1270
+
1271
+ Parameters
1272
+ ----------
1273
+ total_rate : `float`
1274
+ total rate.
1275
+ detectability_condition : `str`
1276
+ detectability condition.
1277
+ param_type : `str`
1278
+ type of parameters.
1279
+ default param_type = 'unlensed'. Other options is 'lensed'.
1280
+ """
937
1281
 
1282
+ data = load_json(self.ler_directory+"/"+self.json_file_names["ler_params"])
938
1283
  # write the results
939
- data["detectable_unlensed_rate_per_year"] = total_rate
940
- data["detectability_condition"] = detectability_condition
1284
+ data[f"detectable_{param_type}_rate_per_year"] = total_rate
1285
+ data[f"detectability_condition+{param_type}"] = detectability_condition
941
1286
  append_json(self.ler_directory+"/"+self.json_file_names["ler_params"], data, replace=True)
942
-
943
- return total_rate, unlensed_param
944
1287
 
945
1288
  def lensed_cbc_statistics(
946
1289
  self, size=None, save_batch=False, resume=False, output_jsonfile=None,
947
1290
  ):
948
1291
  """
949
- Function to generate lensed GW source parameters. This function also stores the parameters in json file.
1292
+ Function to generate lensed GW source parameters. This function calls the lensed_sampling_routine function to generate the parameters in batches. The generated parameters are stored in a json file; and if save_batch=True, it keeps updating the file in batches.
950
1293
 
951
1294
  Parameters
952
1295
  ----------
953
1296
  size : `int`
954
1297
  number of samples.
955
1298
  default size = 100000.
1299
+ save_batch : `bool`
1300
+ if True, the function will save the parameters in batches. if False, the function will save all the parameters at the end of sampling. save_batch=False is faster.
956
1301
  resume : `bool`
957
1302
  resume = False (default) or True.
958
1303
  if True, the function will resume from the last batch.
@@ -963,8 +1308,7 @@ class LeR(LensGalaxyParameterDistribution):
963
1308
  Returns
964
1309
  ----------
965
1310
  lensed_param : `dict`
966
- dictionary of lensed GW source parameters.
967
- lensed_param.keys() =
1311
+ dictionary of lensed GW source parameters. Refer to :attr:`~lensed_param` for details.
968
1312
 
969
1313
  Examples
970
1314
  ----------
@@ -1013,25 +1357,26 @@ class LeR(LensGalaxyParameterDistribution):
1013
1357
 
1014
1358
  def lensed_sampling_routine(self, size, output_jsonfile, save_batch=True, resume=False):
1015
1359
  """
1016
- Function to generate lensed GW source parameters. This function also stores the parameters in json file.
1360
+ Function to generate lensed GW source parameters. This function also stores the parameters in json file in the current batch if save_batch=True.
1017
1361
 
1018
1362
  Parameters
1019
1363
  ----------
1020
1364
  size : `int`
1021
1365
  number of samples.
1022
1366
  default size = 100000.
1023
- resume : `bool`
1024
- resume = False (default) or True.
1025
- if True, the function will resume from the last batch.
1026
1367
  output_jsonfile : `str`
1027
1368
  json file name for storing the parameters.
1028
- default output_jsonfile = 'lensed_params.json'.
1369
+ default output_jsonfile = 'lensed_params.json'. Note that this file will be stored in the self.ler_directory.
1370
+ save_batch : `bool`
1371
+ if True, the function will save the parameters in batches. if False, the function will save all the parameters at the end of sampling. save_batch=False is faster.
1372
+ resume : `bool`
1373
+ resume = False (default) or True.
1374
+ if True, it appends the new samples to the existing json file.
1029
1375
 
1030
1376
  Returns
1031
1377
  ----------
1032
1378
  lensed_param : `dict`
1033
- dictionary of lensed GW source parameters.
1034
- lensed_param.keys() =
1379
+ dictionary of lensed GW source parameters. Refer to :attr:`~lensed_param` for details.
1035
1380
  """
1036
1381
 
1037
1382
  print("sampling lensed params...")
@@ -1046,6 +1391,7 @@ class LeR(LensGalaxyParameterDistribution):
1046
1391
  if len(lensed_param) == 0: # if empty
1047
1392
  lensed_param = lensed_param_
1048
1393
  else:
1394
+ # below will not be used in the first iteration
1049
1395
  # replace the values of the keys
1050
1396
  # idx defines the position that does not satisfy the strong lensing condition
1051
1397
  for key, value in lensed_param_.items():
@@ -1096,15 +1442,21 @@ class LeR(LensGalaxyParameterDistribution):
1096
1442
  self,
1097
1443
  lensed_param=None,
1098
1444
  snr_threshold=[8.0,8.0],
1445
+ pdet_threshold=0.5,
1099
1446
  num_img=[1,1],
1100
1447
  output_jsonfile=None,
1101
1448
  nan_to_num=True,
1102
1449
  detectability_condition="step_function",
1103
1450
  snr_recalculation=False,
1104
- snr_threshold_recalculation=[5.5,5.5],
1451
+ snr_threshold_recalculation=[[4,4], [20,20]],
1105
1452
  ):
1106
1453
  """
1107
- Function to calculate the lensed rate. This function also stores the parameters of the detectable events in json file.
1454
+ Function to calculate the lensed rate. This function also stores the parameters of the detectable events in json file. There are two conditions for detectability: 'step_function' and 'pdet'.
1455
+
1456
+ 1. 'step_function': If two images have SNR>8.0, then the event is detectable. This is a step function. This is with the assumption that SNR function is provided and not None.
1457
+ 2. 'pdet':
1458
+ i) If self.pdet is None and self.snr is not None, then it will calculate the pdet from the snr. There is no hard cut for this pdet and can have value ranging from 0 to 1 near the threshold.
1459
+ ii) If self.pdet is not None, then it will use the generated pdet.
1108
1460
 
1109
1461
  Parameters
1110
1462
  ----------
@@ -1112,11 +1464,14 @@ class LeR(LensGalaxyParameterDistribution):
1112
1464
  dictionary of GW source parameters or json file name.
1113
1465
  default lensed_param = 'lensed_params.json'.
1114
1466
  snr_threshold : `float`
1115
- threshold for detection signal to noise ratio.
1467
+ threshold for detection signal to noise ratio. This is use when self.snr is provided.
1116
1468
  default snr_threshold = [8.0,8.0].
1469
+ pdet_threshold : `float`
1470
+ threshold for detection probability. This is use when self.pdet is provided.
1471
+ default pdet_threshold = 0.5.
1117
1472
  num_img : `int`
1118
- number of images.
1119
- default num_img = [1,1].
1473
+ number of images corresponding to the snr_threshold.
1474
+ default num_img = [1,1]. Together with snr_threshold = [8.0,8.0], it means that two images with snr>8.0. Same condition can also be represented by snr_threshold = 8.0 and num_img = 2.
1120
1475
  output_jsonfile : `str`
1121
1476
  json file name for storing the parameters of the detectable events.
1122
1477
  default output_jsonfile = 'lensed_params_detectable.json'.
@@ -1128,18 +1483,18 @@ class LeR(LensGalaxyParameterDistribution):
1128
1483
  default detectability_condition = 'step_function'.
1129
1484
  other options are 'pdet'.
1130
1485
  snr_recalculation : `bool`
1131
- if True, the SNR of centain events (snr>snr_threshold_recalculation)will be recalculate with 'inner product'. This is useful when the snr is calculated with 'ann' method.
1486
+ if True, the SNR of centain events (snr>snr_threshold_recalculation)will be recalculate with 'inner-product' method. This is useful when the snr is calculated with 'ann' method of `gwsnr`.
1132
1487
  default snr_recalculation = False.
1133
- snr_threshold_recalculation : `float`
1134
- threshold for recalculation of detection signal to noise ratio.
1488
+ snr_threshold_recalculation : `list`
1489
+ lower and upper threshold for recalculation of detection signal to noise ratio.
1490
+ default snr_threshold_recalculation = [[4,4], [20,20]].
1135
1491
 
1136
1492
  Returns
1137
1493
  ----------
1138
1494
  total_rate : `float`
1139
1495
  total lensed rate (Mpc^-3 yr^-1).
1140
1496
  lensed_param : `dict`
1141
- dictionary of lensed GW source parameters of the detectable events.
1142
- lensed_param.keys() =
1497
+ dictionary of lensed GW source parameters of the detectable events. Refer to :attr:`~lensed_param` for details.
1143
1498
 
1144
1499
  Examples
1145
1500
  ----------
@@ -1148,204 +1503,200 @@ class LeR(LensGalaxyParameterDistribution):
1148
1503
  >>> ler.lensed_cbc_statistics();
1149
1504
  >>> total_rate, lensed_param_detectable = ler.lensed_rate()
1150
1505
  """
1151
-
1152
- # call self.json_file_names["ler_param"] and for adding the final results
1153
- data = load_json(self.ler_directory+"/"+self.json_file_names["ler_params"])
1154
1506
 
1155
- # get gw params from json file if not provided
1156
- if lensed_param is None:
1157
- lensed_param = self.json_file_names["lensed_param"]
1158
- if type(lensed_param) == str:
1159
- self.json_file_names["lensed_param"] = lensed_param
1160
- path_ = self.ler_directory+"/"+lensed_param
1161
- print(f"getting lensed_params from json file {path_} ")
1162
- lensed_param = get_param_from_json(self.ler_directory+"/"+lensed_param)
1163
- else:
1164
- print("using provided lensed_param dict...")
1165
- lensed_param = lensed_param.copy()
1507
+ # load lensed parameters
1508
+ lensed_param = self._load_param(lensed_param, param_type="lensed")
1166
1509
 
1167
- # check for images with snr above threshold
1168
- # convert to array
1169
- snr_threshold_ = np.array([snr_threshold]).reshape(-1)
1170
- num_img_ = np.array([num_img]).reshape(-1)
1171
- # get descending sorted idx of snr_threshold
1172
- idx = np.argsort(-snr_threshold_)
1173
- snr_threshold = snr_threshold_[idx]
1174
- num_img = num_img_[idx]
1510
+ # re-analyse the provided snr_threshold and num_img
1511
+ snr_threshold, num_img = self._check_snr_threshold_lensed(snr_threshold, num_img)
1175
1512
 
1176
1513
  # get size of the lensed_param for a parameter
1177
1514
  total_events = len(lensed_param["zs"])
1178
- # recalculate snr if required
1515
+
1179
1516
  # this ensures that the snr is recalculated for the detectable events
1180
1517
  # with inner product
1181
- ############################
1182
- # for ANN SNR #
1183
- ############################
1518
+ # below is use when the snr is calculated with 'ann' method of `gwsnr`
1184
1519
  if snr_recalculation:
1185
- # dealing with provided snr_threshold_recalculation
1186
- snr_threshold_recalculation = np.array([snr_threshold_recalculation]).reshape(-1)
1187
- idx = np.argsort(-snr_threshold_recalculation)
1188
- snr_threshold_recalculation = snr_threshold_recalculation[idx]
1189
- num_img_recalculation = num_img[idx]
1190
-
1191
- # check optimal_snr_net is provided in dict
1192
- snr_param = lensed_param["optimal_snr_net"]
1193
- if "optimal_snr_net" in lensed_param.keys():
1194
- snr_param = -np.sort(-snr_param, axis=1) # sort snr in descending order
1195
- else:
1196
- print("optimal_snr_net not provided in lensed_param dict. Exiting...")
1197
- return None
1198
- # 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
1199
- j = 0
1200
- idx_max = 0
1201
- snr_hit = np.full(total_events, True) # boolean array to store the result of the threshold condition
1202
- for i in range(len(snr_threshold_recalculation)):
1203
- idx_max = idx_max + num_img_recalculation[i]
1204
- snr_hit = snr_hit & (np.sum((snr_param[:,j:idx_max] > snr_threshold_recalculation[i]), axis=1) >= num_img_recalculation[i])
1205
- j = idx_max
1206
-
1207
- # reduce the size of the dict
1208
- for key, value in lensed_param.items():
1209
- lensed_param[key] = value[snr_hit]
1210
- # recalculate more accurate snrs
1211
- print("calculating snrs...")
1212
- snrs, lensed_param = self.get_lensed_snrs(
1213
- snr_calculator=self.snr_bilby,
1214
- list_of_detectors=self.list_of_detectors,
1215
- lensed_param=lensed_param,
1216
- )
1217
- lensed_param.update(snrs)
1218
- ############################
1219
-
1220
- if self.snr:
1221
- if "optimal_snr_net" not in lensed_param:
1222
- raise ValueError("'optimal_snr_net' not in lensed parm dict provided")
1223
- if detectability_condition == "step_function":
1224
- print("given detectability_condition == 'step_function'")
1225
-
1226
- if "optimal_snr_net" in lensed_param.keys():
1227
- snr_param = lensed_param["optimal_snr_net"]
1228
- snr_param = -np.sort(-snr_param, axis=1) # sort snr in descending order
1229
- else:
1230
- print("snr not provided in lensed_param dict. Exiting...")
1231
- return None
1232
- snr_hit = np.full(len(snr_param), True) # boolean array to store the result of the threshold condition
1233
-
1234
- # 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
1235
- j = 0
1236
- idx_max = 0
1237
- for i in range(len(snr_threshold)):
1238
- idx_max = idx_max + num_img[i]
1239
- snr_hit = snr_hit & (np.sum((snr_param[:,j:idx_max] > snr_threshold[i]), axis=1) >= num_img[i])
1240
- j = idx_max
1241
-
1242
- # montecarlo integration
1243
- total_rate = self.normalization_pdf_z_lensed * np.sum(snr_hit)/total_events
1244
- print("total lensed rate (yr^-1) (with step function): {}".format(total_rate))
1245
-
1246
- elif detectability_condition == "pdet":
1247
- print("given detectability_condition == 'pdet'")
1248
- # pdet dimension is (size, n_max_images)
1249
- snr_param = lensed_param["optimal_snr_net"]
1250
- snr_param = -np.sort(-snr_param, axis=1) # sort snr in descending order
1251
-
1252
- pdet = np.ones(np.shape(snr_param))
1253
- j = 0
1254
- idx_max = 0
1255
- for i in range(len(snr_threshold)):
1256
- idx_max = idx_max + num_img[i]
1257
- pdet[:,j:idx_max] = 1 - norm.cdf(snr_threshold[i] - snr_param[:,j:idx_max])
1258
- j = idx_max
1259
-
1260
- snr_hit = np.prod(pdet, axis=1)>0.5
1261
- # montecarlo integration
1262
- total_rate = self.normalization_pdf_z_lensed * np.sum(snr_hit)/total_events
1263
- print(f"total lensed rate (yr^-1) (with pdet function): {total_rate}")
1264
-
1265
- elif self.pdet:
1266
- # check if pdet is provided in lensed_param dict
1267
- if "pdet_net" in lensed_param.keys():
1268
- print("given detectability_condition == 'pdet'")
1269
- pdet = lensed_param["pdet_net"]
1270
- pdet = -np.sort(-pdet, axis=1) # sort pdet in descending order
1271
- pdet = pdet[:,:np.sum(num_img)] # use only num_img images
1520
+ lensed_param = self._recalculate_snr_lensed(lensed_param, snr_threshold_recalculation, num_img, total_events)
1272
1521
 
1273
- snr_hit = np.prod(pdet, axis=1)>0.5
1274
- # montecarlo integration
1275
- total_rate = self.normalization_pdf_z_lensed * np.sum(snr_hit)/total_events
1276
- print(f"total lensed rate (yr^-1) (with pdet function): {total_rate}")
1522
+ snr_hit = self._find_detectable_index_lensed(lensed_param, snr_threshold, pdet_threshold, num_img, detectability_condition)
1277
1523
 
1278
- else:
1279
- raise ValueError("'pdet_net' not in unlensed parm dict provided")
1524
+ # montecarlo integration
1525
+ total_rate = self.rate_function(np.sum(snr_hit), total_events, param_type="lensed")
1280
1526
 
1281
- print(f"number of simulated lensed detectable events: {np.sum(snr_hit)}")
1282
- print(f"number of simulated all lensed events: {total_events}")
1283
- # store all detectable params in json file
1284
- if nan_to_num:
1285
- for key, value in lensed_param.items():
1286
- lensed_param[key] = np.nan_to_num(value[snr_hit])
1287
- else:
1288
- for key, value in lensed_param.items():
1289
- lensed_param[key] = value[snr_hit]
1290
-
1291
1527
  # store all detectable params in json file
1292
- if output_jsonfile is None:
1293
- output_jsonfile = self.json_file_names["lensed_param_detectable"]
1294
- else:
1295
- self.json_file_names["lensed_param_detectable"] = output_jsonfile
1296
- path_ = self.ler_directory+"/"+output_jsonfile
1297
- print(f"storing detectable lensed params in {path_}")
1298
- append_json(self.ler_directory+"/"+output_jsonfile, lensed_param, replace=True)
1528
+ self._save_detectable_params(output_jsonfile, lensed_param, snr_hit, key_file_name="lensed_param_detectable", nan_to_num=nan_to_num, verbose=True, replace_jsonfile=True)
1299
1529
 
1300
- # write the results
1301
- data["detectable_lensed_rate_per_year"] = total_rate
1302
- data["detectability_condition"] = detectability_condition
1303
- append_json(self.ler_directory+"/"+self.json_file_names["ler_params"], data, replace=True)
1530
+ # append ler_param and save it
1531
+ self._append_ler_param(total_rate, detectability_condition, param_type="lensed")
1304
1532
 
1305
1533
  return total_rate, lensed_param
1306
-
1307
- def rate_ratio(self):
1534
+
1535
+ def _check_snr_threshold_lensed(self, snr_threshold, num_img):
1308
1536
  """
1309
- Function to calculate and display unlensed and lensed merger rate ratio.
1310
- It will get the unlensed_rate and lensed_rate from self.json_file_ler_param
1537
+ Helper function to check the snr_threshold and num_img for lensed events.
1311
1538
 
1312
- Returns
1539
+ Parameters
1313
1540
  ----------
1314
- rate_ratio : `float`
1315
- rate ratio.
1541
+ snr_threshold : `float`
1542
+ threshold for detection signal to noise ratio.
1543
+ default snr_threshold = [8.0,8.0].
1544
+ num_img : `int`
1545
+ number of images corresponding to the snr_threshold.
1546
+ default num_img = [1,1]. Together with snr_threshold = [8.0,8.0], it means that two images with snr>8.0. Same condition can also be represented by snr_threshold = 8.0 and num_img = 2.
1316
1547
 
1317
- Examples
1548
+ Returns
1318
1549
  ----------
1319
- >>> from ler.rates import LeR
1320
- >>> ler = LeR()
1321
- >>> ler.unlensed_cbc_statistics();
1322
- >>> ler.lensed_cbc_statistics();
1323
- >>> ler.unlensed_rate();
1324
- >>> ler.lensed_rate();
1325
- >>> ler.rate_ratio()
1550
+ snr_threshold : `float`
1551
+ threshold for detection signal to noise ratio.
1552
+ num_img : `int`
1553
+ number of images corresponding to the snr_threshold.
1326
1554
  """
1555
+ # check for images with snr above threshold
1556
+ # convert to array
1557
+ snr_threshold_ = np.array([snr_threshold]).reshape(-1)
1558
+ num_img_ = np.array([num_img]).reshape(-1)
1559
+ # get descending sorted idx of snr_threshold
1560
+ idx = np.argsort(-snr_threshold_)
1561
+ snr_threshold = snr_threshold_[idx]
1562
+ num_img = num_img_[idx]
1327
1563
 
1328
- # call json_file_ler_param and add the results
1329
- data = load_json(self.ler_directory+"/"+self.json_file_names["ler_params"])
1330
-
1331
- try:
1332
- unlensed_rate = data['detectable_unlensed_rate_per_year']
1333
- lensed_rate = data['detectable_lensed_rate_per_year']
1334
- except:
1335
- raise ValueError(f"detectable_unlensed_rate_per_year or 'detectable_lensed_rate_per_year' not found in {self.json_file_names['ler_params']} json file. Exiting...")
1336
-
1337
- rate_ratio = unlensed_rate / lensed_rate
1338
- # append the results
1339
- data['rate_ratio'] = rate_ratio
1340
- # write the results
1341
- append_json(self.ler_directory+"/"+self.json_file_names["ler_params"], data, replace=True)
1342
-
1343
- print(f"unlensed_rate: {unlensed_rate}")
1344
- print(f"lensed_rate: {lensed_rate}")
1345
- print(f"ratio: {rate_ratio}")
1564
+ return snr_threshold, num_img
1565
+
1566
+ def _recalculate_snr_lensed(self, lensed_param, snr_threshold_recalculation, num_img, total_events):
1567
+ """
1568
+ Helper function to recalculate the SNR of lensed events with SNR above a given threshold.
1569
+
1570
+ Parameters
1571
+ ----------
1572
+ lensed_param : `dict`
1573
+ dictionary of lensed GW source parameters.
1574
+ snr_threshold_recalculation : `list`
1575
+ lower and upper threshold for recalculation of detection signal to noise ratio.
1576
+ num_img : `int`
1577
+ number of images corresponding to the snr_threshold.
1578
+ total_events : `int`
1579
+ total number of events.
1580
+
1581
+ Returns
1582
+ ----------
1583
+ lensed_param : `dict`
1584
+ dictionary of lensed GW source parameters.
1585
+ """
1586
+
1587
+ # dealing with provided snr_threshold_recalculation
1588
+ snr_threshold_recalculation_min, num_img_recalculation = self._check_snr_threshold_lensed(snr_threshold_recalculation[0], num_img)
1589
+ snr_threshold_recalculation_max, _ = self._check_snr_threshold_lensed(snr_threshold_recalculation[1], num_img)
1590
+
1591
+ # check optimal_snr_net is provided in dict
1592
+ if "optimal_snr_net" not in lensed_param:
1593
+ raise ValueError("optimal_snr_net not provided in lensed_param dict. Exiting...")
1594
+
1595
+ snr_param = lensed_param["optimal_snr_net"]
1596
+ snr_param = -np.sort(-snr_param, axis=1) # sort snr in descending order
1597
+
1598
+ # 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
1599
+ j = 0
1600
+ idx_max = 0
1601
+ snr_hit = np.full(total_events, True) # boolean array to store the result of the threshold condition
1602
+ for i, snr_th_min in enumerate(snr_threshold_recalculation_min):
1603
+ snr_th_max = snr_threshold_recalculation_max[i]
1604
+ idx_max = idx_max + num_img_recalculation[i]
1605
+ snr_hit = snr_hit & (np.sum((snr_param[:,j:idx_max] > snr_th_min) & (snr_param[:,j:idx_max] < snr_th_max), axis=1) >= num_img_recalculation[i])
1606
+ j = idx_max
1607
+
1608
+ # reduce the size of the dict
1609
+ for key, value in lensed_param.items():
1610
+ lensed_param[key] = value[snr_hit]
1611
+ # recalculate more accurate snrs
1612
+ print("calculating snrs...")
1613
+ snrs, lensed_param = self.get_lensed_snrs(
1614
+ snr_calculator=self.snr_bilby,
1615
+ list_of_detectors=self.list_of_detectors,
1616
+ lensed_param=lensed_param,
1617
+ )
1618
+ lensed_param.update(snrs)
1619
+
1620
+ return lensed_param
1621
+
1622
+ def _find_detectable_index_lensed(self, lensed_param, snr_threshold, pdet_threshold, num_img, detectability_condition):
1623
+ """
1624
+ Helper function to find the index of detectable events based on SNR or p_det.
1625
+
1626
+ Parameters
1627
+ ----------
1628
+ lensed_param : `dict`
1629
+ dictionary of lensed GW source parameters.
1630
+ snr_threshold : `float` or `list`
1631
+ threshold for detection signal to noise ratio.
1632
+ default snr_threshold = [8.0,8.0].
1633
+ pdet_threshold : `float` or `list`
1634
+ threshold for detection probability.
1635
+ default pdet_threshold = 0.5.
1636
+ num_img : `int` or `list`
1637
+ number of images corresponding to the snr_threshold.
1638
+ default num_img = [1,1].
1639
+ detectability_condition : `str`
1640
+ detectability condition.
1641
+ default detectability_condition = 'step_function'.
1642
+ other options are 'pdet'.
1643
+
1644
+ Returns
1645
+ ----------
1646
+ snr_hit : `bool`
1647
+ boolean array to store the result of the threshold condition.
1648
+ """
1649
+
1650
+ print(f"given detectability_condition == {detectability_condition}")
1651
+ if detectability_condition == "step_function":
1652
+ if "optimal_snr_net" not in lensed_param:
1653
+ raise ValueError("'optimal_snr_net' not in lensed parm dict provided")
1654
+ snr_param = lensed_param["optimal_snr_net"]
1655
+ snr_param = -np.sort(-snr_param, axis=1) # sort snr in descending order
1656
+ snr_hit = np.full(len(snr_param), True) # boolean array to store the result of the threshold condition
1657
+
1658
+ # 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
1659
+ # algorithm:
1660
+ # 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.
1661
+ # ii) for loop runs wrt snr_threshold. idx_max = idx_max + num_img[i]
1662
+ # 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.
1663
+ # 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.
1664
+ j = 0
1665
+ idx_max = 0
1666
+ for i, snr_th in enumerate(snr_threshold):
1667
+ idx_max = idx_max + num_img[i]
1668
+ snr_hit = snr_hit & (np.sum((snr_param[:,j:idx_max] > snr_th), axis=1) >= num_img[i])
1669
+ j = idx_max
1670
+
1671
+ elif detectability_condition == "pdet":
1672
+ if "pdet_net" not in lensed_param:
1673
+ if "optimal_snr_net" not in lensed_param:
1674
+ raise ValueError("'optimal_snr_net' or 'pdet_net' not in lensed parm dict provided")
1675
+ else:
1676
+ print("calculating pdet using 'optimal_snr_net'...")
1677
+ # pdet dimension is (size, n_max_images)
1678
+ snr_param = lensed_param["optimal_snr_net"]
1679
+ snr_param = -np.sort(-snr_param, axis=1) # sort snr in descending order
1680
+
1681
+ # column index beyong np.sum(num_img)-1 are not considered
1682
+ pdet = np.ones((len(snr_param), np.sum(num_img)))
1683
+ j = 0
1684
+ idx_max = 0
1685
+ for i, snr_th in enumerate(snr_threshold):
1686
+ idx_max = idx_max + num_img[i]
1687
+ pdet[:,j:idx_max] = 1 - norm.cdf(snr_th - snr_param[:,j:idx_max])
1688
+ j = idx_max
1689
+ else:
1690
+ pdet = lensed_param["pdet_net"]
1691
+ # sort pdet in descending order
1692
+ pdet = -np.sort(-pdet, axis=1)
1693
+ # column index beyong np.sum(num_img)-1 are not considered
1694
+ pdet = pdet[:,:np.sum(num_img)]
1695
+
1696
+ snr_hit = np.prod(pdet, axis=1)>=pdet_threshold
1697
+
1698
+ return snr_hit
1346
1699
 
1347
- return unlensed_rate / lensed_rate
1348
-
1349
1700
  def rate_comparision_with_rate_calculation(
1350
1701
  self,
1351
1702
  unlensed_param=None,
@@ -1359,7 +1710,7 @@ class LeR(LensGalaxyParameterDistribution):
1359
1710
  detectability_condition="step_function",
1360
1711
  ):
1361
1712
  """
1362
- Function to calculate the unlensed and lensed rate and compare by computing the ratio. This function also stores the parameters of the detectable events in json file.
1713
+ Function to calculate the unlensed and lensed rate and compare by computing the ratio. This function also stores the parameters of the detectable events in json file. If you use this function, you do not need to call the functions unlensed_rate and lensed_rate separately.
1363
1714
 
1364
1715
  Parameters
1365
1716
  ----------
@@ -1368,21 +1719,27 @@ class LeR(LensGalaxyParameterDistribution):
1368
1719
  default unlensed_param = 'unlensed_params.json'.
1369
1720
  snr_threshold_unlensed : `float`
1370
1721
  threshold for detection signal to noise ratio.
1371
- e.g. snr_threshold = 8.
1722
+ e.g. snr_threshold_unlensed = 8.
1372
1723
  output_jsonfile_unlensed : `str`
1373
1724
  json file name for storing the parameters of the detectable events.
1374
- default output_jsonfile = 'unlensed_params_detectable.json'.
1725
+ default output_jsonfile_unlensed = 'unlensed_params_detectable.json'.
1375
1726
  lensed_param : `dict` or `str`
1376
1727
  dictionary of GW source parameters or json file name.
1377
1728
  default lensed_param = 'lensed_params.json'.
1378
1729
  snr_threshold_lensed : `float`
1379
1730
  threshold for detection signal to noise ratio.
1380
- e.g. snr_threshold = 8.
1731
+ default snr_threshold_lensed = [8.0,8.0].
1732
+ num_img : `int`
1733
+ number of images.
1734
+ default num_img = [1,1]. Together with snr_threshold = [8.0,8.0], it means that two images with snr>8.0. Same condition can also be represented by snr_threshold = 8.0 and num_img = 2.
1381
1735
  output_jsonfile_lensed : `str`
1382
1736
  json file name for storing the parameters of the detectable events.
1383
- default output_jsonfile = 'lensed_params_detectable.json'.
1737
+ default output_jsonfile_lensed = 'lensed_params_detectable.json'.
1738
+ nan_to_num : `bool`
1739
+ if True, nan values will be converted to 0.
1740
+ default nan_to_num = True.
1384
1741
  detectability_condition : `str`
1385
- detectability condition.
1742
+ detectability condition.
1386
1743
  default detectability_condition = 'step_function'.
1387
1744
  other options are 'pdet'.
1388
1745
 
@@ -1391,9 +1748,9 @@ class LeR(LensGalaxyParameterDistribution):
1391
1748
  rate_ratio : `float`
1392
1749
  rate ratio.
1393
1750
  unlensed_param : `dict`
1394
- dictionary of unlensed GW source parameters of the detectable events.
1751
+ dictionary of unlensed GW source parameters of the detectable events. Refer to :attr:`~unlensed_param` for details.
1395
1752
  lensed_param : `dict`
1396
- dictionary of lensed GW source parameters of the detectable events.
1753
+ dictionary of lensed GW source parameters of the detectable events. Refer to :attr:`~lensed_param` for details.
1397
1754
 
1398
1755
  Examples
1399
1756
  ----------
@@ -1424,200 +1781,169 @@ class LeR(LensGalaxyParameterDistribution):
1424
1781
  detectability_condition=detectability_condition,
1425
1782
  )
1426
1783
  # calculate rate ratio
1784
+ rate_ratio = self.rate_ratio()
1785
+
1786
+ return rate_ratio, unlensed_param, lensed_param
1787
+
1788
+ def rate_ratio(self):
1789
+ """
1790
+ Function to calculate and display unlensed and lensed merger rate ratio. It will get the unlensed_rate and lensed_rate from files corresponding to the names included in self.json_file_ler_param.
1791
+
1792
+ Returns
1793
+ ----------
1794
+ rate_ratio : `float`
1795
+ rate ratio.
1796
+
1797
+ Examples
1798
+ ----------
1799
+ >>> from ler.rates import LeR
1800
+ >>> ler = LeR()
1801
+ >>> ler.unlensed_cbc_statistics();
1802
+ >>> ler.lensed_cbc_statistics();
1803
+ >>> ler.unlensed_rate();
1804
+ >>> ler.lensed_rate();
1805
+ >>> ler.rate_ratio()
1806
+ """
1807
+
1808
+ # call json_file_ler_param and add the results
1809
+ data = load_json(self.ler_directory+"/"+self.json_file_names["ler_params"])
1810
+
1811
+ try:
1812
+ unlensed_rate = data['detectable_unlensed_rate_per_year']
1813
+ lensed_rate = data['detectable_lensed_rate_per_year']
1814
+ except:
1815
+ raise ValueError(f"detectable_unlensed_rate_per_year or 'detectable_lensed_rate_per_year' not found in {self.json_file_names['ler_params']} json file. Exiting...")
1816
+
1427
1817
  rate_ratio = unlensed_rate / lensed_rate
1428
1818
  # append the results
1429
- data['detectable_unlensed_rate_per_year'] = unlensed_rate
1430
- data['detectable_lensed_rate_per_year'] = lensed_rate
1431
1819
  data['rate_ratio'] = rate_ratio
1432
- data['detectability_condition'] = detectability_condition
1433
1820
  # write the results
1434
1821
  append_json(self.ler_directory+"/"+self.json_file_names["ler_params"], data, replace=True)
1435
1822
 
1436
- print(f"unlensed_rate (per year): {unlensed_rate}")
1437
- print(f"lensed_rate (per year): {lensed_rate}")
1823
+ print(f"unlensed_rate: {unlensed_rate}")
1824
+ print(f"lensed_rate: {lensed_rate}")
1438
1825
  print(f"ratio: {rate_ratio}")
1439
1826
 
1440
- return rate_ratio, unlensed_param, lensed_param
1827
+ return unlensed_rate / lensed_rate
1441
1828
 
1442
1829
  def selecting_n_unlensed_detectable_events(
1443
1830
  self,
1444
1831
  size=100,
1445
1832
  batch_size=None,
1446
1833
  snr_threshold=8.0,
1834
+ pdet_threshold=0.5,
1447
1835
  resume=False,
1448
1836
  output_jsonfile="n_unlensed_param_detectable.json",
1449
1837
  meta_data_file="meta_unlensed.json",
1450
1838
  detectability_condition="step_function",
1451
1839
  trim_to_size=True,
1452
1840
  snr_recalculation=False,
1453
- snr_threshold_recalculation=5.5,
1841
+ snr_threshold_recalculation=[4, 12],
1454
1842
  ):
1455
1843
  """
1456
- Function to select n unlensed detectable events.
1844
+ Function to generate n unlensed detectable events. This fuction samples the unlensed parameters and save only the detectable events in json file. It also records metadata in the JSON file, which includes the total number of events and the cumulative rate of events. This functionality is particularly useful for generating a fixed or large number of detectable events until the event rates stabilize.
1457
1845
 
1458
1846
  Parameters
1459
1847
  ----------
1460
1848
  size : `int`
1461
1849
  number of samples to be selected.
1462
1850
  default size = 100.
1851
+ batch_size : `int`
1852
+ batch size for sampling.
1853
+ default batch_size = 50000.
1463
1854
  snr_threshold : `float`
1464
1855
  threshold for detection signal to noise ratio.
1465
1856
  e.g. snr_threshold = 8.
1857
+ pdet_threshold : `float`
1858
+ threshold for detection probability.
1859
+ default pdet_threshold = 0.5.
1466
1860
  resume : `bool`
1467
- if True, it will resume the sampling from the last batch.
1468
- default resume = False.
1861
+ resume = False (default) or True.
1862
+ if True, the function will resume from the last batch.
1469
1863
  output_jsonfile : `str`
1470
- json file name for storing the parameters.
1471
- default output_jsonfile = 'n_unlensed_params_detectable.json'.
1864
+ json file name for storing the parameters of the detectable events.
1865
+ default output_jsonfile = 'n_unlensed_param_detectable.json'.
1866
+ meta_data_file : `str`
1867
+ json file name for storing the metadata.
1868
+ default meta_data_file = 'meta_unlensed.json'.
1869
+ detectability_condition : `str`
1870
+ detectability condition.
1871
+ default detectability_condition = 'step_function'.
1872
+ other options are 'pdet'.
1873
+ trim_to_size : `bool`
1874
+ if True, the final result will be trimmed to size.
1875
+ default trim_to_size = True.
1876
+ snr_recalculation : `bool`
1877
+ if True, the SNR of centain events (snr>snr_threshold_recalculation)will be recalculate with 'inner-product' method. This is useful when the snr is calculated with 'ann' method of `gwsnr`.
1878
+ default snr_recalculation = False.
1879
+ snr_threshold_recalculation : `list`
1880
+ lower and upper threshold for recalculation of detection signal to noise ratio.
1881
+ default snr_threshold_recalculation = [4, 12].
1472
1882
 
1473
1883
  Returns
1474
1884
  ----------
1475
1885
  param_final : `dict`
1476
- dictionary of unlensed GW source parameters of the detectable events.
1477
- param_final.keys() = ['zs', 'geocent_time', 'ra', 'dec', 'phase', 'psi', 'theta_jn', 'luminosity_distance', 'mass_1_source', 'mass_2_source', 'mass_1', 'mass_2', 'optimal_snr_net', 'L1', 'H1', 'V1']
1886
+ dictionary of unlensed GW source parameters of the detectable events. Refer to :attr:`~unlensed_param` for details.
1478
1887
 
1479
1888
  Examples
1480
1889
  ----------
1481
1890
  >>> from ler.rates import LeR
1482
1891
  >>> ler = LeR()
1483
- >>> unlensed_param_final = ler.selecting_n_unlensed_detectable_events(size=500)
1892
+ >>> unlensed_param = ler.selecting_n_unlensed_detectable_events(size=100)
1484
1893
  """
1485
1894
 
1486
- meta_data_path = self.ler_directory+"/"+meta_data_file
1487
- output_path = self.ler_directory+"/"+output_jsonfile
1488
-
1489
- if batch_size is None:
1490
- batch_size = self.batch_size
1491
- else:
1492
- self.batch_size = batch_size
1895
+ # initial setup
1896
+ 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)
1493
1897
 
1494
- if not resume:
1495
- n = 0 # iterator
1496
- events_total = 0
1497
- if os.path.exists(output_path):
1498
- os.remove(output_path)
1499
- if os.path.exists(meta_data_path):
1500
- os.remove(meta_data_path)
1501
- else:
1502
- # get sample size as size from json file
1503
- if os.path.exists(output_path):
1504
- param_final = get_param_from_json(output_path)
1505
- n = len(param_final["zs"])
1506
- events_total = load_json(meta_data_path)["events_total"][-1]
1507
- del param_final
1508
- else:
1509
- n = 0
1510
- events_total = 0
1511
-
1512
- buffer_file = self.ler_directory+"/"+"unlensed_params_buffer.json"
1513
- print("collected number of detectable events = ", n)
1514
-
1898
+ # loop until n samples are collected
1515
1899
  while n < size:
1516
1900
  # disable print statements
1517
1901
  with contextlib.redirect_stdout(None):
1518
- self.dict_buffer = None
1902
+ self.dict_buffer = None # this is used to store the sampled unlensed_param in batches when running the sampling_routine
1519
1903
  unlensed_param = self.unlensed_sampling_routine(
1520
1904
  size=batch_size, output_jsonfile=buffer_file, save_batch=False,resume=False
1521
1905
  )
1522
1906
 
1523
1907
  total_events_in_this_iteration = len(unlensed_param["zs"])
1524
- ############################
1525
- # for ANN SNR #
1526
- ############################
1908
+ # below is use when the snr is calculated with 'ann' method of `gwsnr`
1527
1909
  if snr_recalculation:
1528
1910
  # select only above centain snr threshold
1529
- param = unlensed_param["optimal_snr_net"]
1530
- idx_detectable = param > snr_threshold_recalculation
1531
- # reduce the size of the dict
1532
- for key, value in unlensed_param.items():
1533
- unlensed_param[key] = value[idx_detectable]
1534
- # recalculate more accurate snrs
1535
- snrs = self.snr_bilby(gw_param_dict=unlensed_param)
1536
- unlensed_param.update(snrs)
1537
- ############################
1538
-
1539
- if self.snr:
1540
- if "optimal_snr_net" not in unlensed_param:
1541
- raise ValueError("'optimal_snr_net' not in unlensed parm dict provided")
1542
- if detectability_condition == "step_function":
1543
- #print("given detectability_condition == 'step_function'")
1544
- param = unlensed_param["optimal_snr_net"]
1545
- threshold = snr_threshold
1546
- elif detectability_condition == "pdet":
1547
- # print("given detectability_condition == 'pdet'")
1548
- param = 1 - norm.cdf(snr_threshold - unlensed_param["optimal_snr_net"])
1549
- unlensed_param["pdet_net"] = param
1550
- threshold = 0.5
1551
- elif self.pdet:
1552
- if "pdet_net" in unlensed_param:
1553
- #print("given detectability_condition == 'pdet'")
1554
- param = unlensed_param["pdet_net"]
1555
- threshold = 0.5
1556
- else:
1557
- raise ValueError("'pdet_net' not in unlensed parm dict provided")
1911
+ unlensed_param = self._recalculate_snr_unlensed(unlensed_param, snr_threshold_recalculation)
1558
1912
 
1559
- # index of detectable events
1560
- idx = param > threshold
1913
+ # find index of detectable events
1914
+ idx_detectable = self._find_detectable_index_unlensed(unlensed_param, snr_threshold, pdet_threshold, detectability_condition)
1561
1915
 
1562
1916
  # store all params in json file
1563
- for key, value in unlensed_param.items():
1564
- unlensed_param[key] = value[idx]
1565
- append_json(file_name=output_path, new_dictionary=unlensed_param, replace=False)
1917
+ self._save_detectable_params(output_jsonfile, unlensed_param, idx_detectable, key_file_name="n_unlensed_detectable_events", nan_to_num=False, verbose=False, replace_jsonfile=False)
1566
1918
 
1567
- n += np.sum(idx)
1919
+ n += np.sum(idx_detectable)
1568
1920
  events_total += total_events_in_this_iteration
1569
- total_rate = self.normalization_pdf_z * n / events_total
1921
+ total_rate = self.rate_function(n, events_total, param_type="unlensed", verbose=False)
1570
1922
 
1571
- # save meta data
1572
- meta_data = dict(events_total=[events_total], detectable_events=[float(n)], total_rate=[total_rate])
1573
- if os.path.exists(meta_data_path):
1574
- append_json(file_name=meta_data_path, new_dictionary=meta_data, replace=False)
1575
- else:
1576
- append_json(file_name=meta_data_path, new_dictionary=meta_data, replace=True)
1577
-
1578
- print("collected number of detectable events = ", n)
1579
- print("total number of events = ", events_total)
1580
- print(f"total unlensed rate (yr^-1): {total_rate}")
1923
+ # bookmark
1924
+ self._append_meta_data(meta_data_path, n, events_total, total_rate)
1581
1925
 
1582
- print(f"storing detectable unlensed params in {output_path}")
1583
- print(f"storing meta data in {meta_data_path}")
1926
+ print(f"stored detectable unlensed params in {output_path}")
1927
+ print(f"stored meta data in {meta_data_path}")
1584
1928
 
1585
1929
  if trim_to_size:
1586
- # trim the final param dictionary
1587
- print(f"\n trmming final result to size={size}")
1588
- param_final = get_param_from_json(output_path)
1589
- # randomly select size number of samples
1590
- len_ = len(list(param_final.values())[0])
1591
- idx = np.random.choice(len_, size, replace=False)
1592
- # trim the final param dictionary, randomly, without repeating
1593
- for key, value in param_final.items():
1594
- param_final[key] = value[idx]
1595
-
1596
- # change meta data
1597
- meta_data = load_json(meta_data_path)
1598
- old_events_total = meta_data["events_total"][-1]
1599
- old_detectable_events = meta_data["detectable_events"][-1]
1600
-
1601
- # adjust the meta data
1602
- # following is to keep rate the same
1603
- new_events_total = np.round(size*old_events_total/old_detectable_events)
1604
- new_total_rate = self.normalization_pdf_z * size / new_events_total
1605
- meta_data["events_total"][-1] = new_events_total
1606
- meta_data["detectable_events"][-1] = size
1607
- meta_data["total_rate"][-1] = new_total_rate
1608
-
1609
-
1610
- print("collected number of detectable events = ", size)
1611
- print("total number of events = ", new_events_total)
1612
- print(f"total unlensed rate (yr^-1): {new_total_rate}")
1613
-
1614
- # save the meta data
1615
- append_json(meta_data_path, meta_data, replace=True)
1616
- # save the final param dictionary
1617
- append_json(output_path, param_final, replace=True)
1930
+ param_final, total_rate = self._trim_results_to_size(size, output_path, meta_data_path)
1618
1931
  else:
1619
1932
  param_final = get_param_from_json(output_path)
1620
1933
 
1934
+ # call self.json_file_names["ler_param"] and for adding the final results
1935
+ data = load_json(self.ler_directory+"/"+self.json_file_names["ler_params"])
1936
+ # write the results
1937
+ try:
1938
+ data["detectable_unlensed_rate_per_year"] = total_rate
1939
+ data["detectability_condition_unlensed"] = detectability_condition
1940
+ except:
1941
+ meta = get_param_from_json(meta_data_path)
1942
+ data["detectable_unlensed_rate_per_year"] = meta["total_rate"][-1]
1943
+ data["detectability_condition_unlensed"] = detectability_condition
1944
+
1945
+ append_json(self.ler_directory+"/"+self.json_file_names["ler_params"], data, replace=True)
1946
+
1621
1947
  return param_final
1622
1948
 
1623
1949
  def selecting_n_lensed_detectable_events(
@@ -1625,6 +1951,7 @@ class LeR(LensGalaxyParameterDistribution):
1625
1951
  size=100,
1626
1952
  batch_size=None,
1627
1953
  snr_threshold=[8.0,8.0],
1954
+ pdet_threshold=0.5,
1628
1955
  num_img=[1,1],
1629
1956
  resume=False,
1630
1957
  detectability_condition="step_function",
@@ -1633,50 +1960,156 @@ class LeR(LensGalaxyParameterDistribution):
1633
1960
  trim_to_size=True,
1634
1961
  nan_to_num=False,
1635
1962
  snr_recalculation=False,
1636
- snr_threshold_recalculation=[5.5,5.5],
1963
+ snr_threshold_recalculation=[[4,4],[12,12]],
1637
1964
  ):
1638
1965
  """
1639
- Function to select n lensed detectable events.
1966
+ Function to generate n lensed detectable events. This fuction only samples the lensed parameters and save only the detectable events in json file. It also records metadata in the JSON file, which includes the total number of events and the cumulative rate of events. This functionality is particularly useful for generating a fixed or large number of detectable events until the event rates stabilize.
1640
1967
 
1641
1968
  Parameters
1642
1969
  ----------
1643
1970
  size : `int`
1644
- number of samples to be selected.
1971
+ number of samples.
1645
1972
  default size = 100.
1973
+ batch_size : `int`
1974
+ batch size for sampling.
1975
+ default batch_size = 50000.
1646
1976
  snr_threshold : `float`
1647
1977
  threshold for detection signal to noise ratio.
1648
- e.g. snr_threshold = 8.
1978
+ default snr_threshold = [8.0,8.0].
1979
+ pdet_threshold : `float`
1980
+ threshold for detection probability.
1981
+ default pdet_threshold = 0.5.
1649
1982
  num_img : `int`
1650
1983
  number of images.
1651
- default num_img = 2.
1984
+ default num_img = [1,1]. Together with snr_threshold = [8.0,8.0], it means that two images with snr>8.0. Same condition can also be represented by snr_threshold = 8.0 and num_img = 2.
1652
1985
  resume : `bool`
1653
- if True, it will resume the sampling from the last batch.
1654
- default resume = False.
1986
+ resume = False (default) or True.
1987
+ if True, it appends the new samples to the existing json file.
1655
1988
  detectability_condition : `str`
1656
1989
  detectability condition.
1657
1990
  default detectability_condition = 'step_function'.
1658
1991
  other options are 'pdet'.
1659
1992
  output_jsonfile : `str`
1660
- json file name for storing the parameters.
1993
+ json file name for storing the parameters of the detectable events.
1994
+ default output_jsonfile = 'n_lensed_params_detectable.json'.
1995
+ meta_data_file : `str`
1996
+ json file name for storing the metadata.
1997
+ default meta_data_file = 'meta_lensed.json'.
1998
+ trim_to_size : `bool`
1999
+ if True, the final result will be trimmed to size.
2000
+ default trim_to_size = True.
2001
+ nan_to_num : `bool`
2002
+ if True, nan values will be converted to 0.
2003
+ default nan_to_num = False.
2004
+ snr_recalculation : `bool`
2005
+ if True, the SNR of centain events (snr>snr_threshold_recalculation)will be recalculate with 'inner-product' method. This is useful when the snr is calculated with 'ann' method of `gwsnr`.
2006
+ default snr_recalculation = False.
2007
+ snr_threshold_recalculation : `list`
2008
+ lower and upper threshold for recalculation of detection signal to noise ratio.
2009
+ default snr_threshold_recalculation = [[4,4], [12,12]].
1661
2010
 
1662
2011
  Returns
1663
2012
  ----------
1664
2013
  param_final : `dict`
1665
- dictionary of lensed GW source parameters of the detectable events.
1666
- param_final.keys() =
2014
+ dictionary of lensed GW source parameters of the detectable events. Refer to :attr:`~lensed_param` for details.
1667
2015
 
1668
2016
  Examples
1669
2017
  ----------
1670
2018
  >>> from ler.rates import LeR
1671
2019
  >>> ler = LeR()
1672
- >>> lensed_param_final = ler.selecting_n_lensed_detectable_events(size=500)
2020
+ >>> lensed_param = ler.selecting_n_lensed_detectable_events(size=100)
2021
+ """
2022
+
2023
+ # initial setup
2024
+ 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)
2025
+
2026
+ # re-analyse the provided snr_threshold and num_img
2027
+ snr_threshold, num_img = self._check_snr_threshold_lensed(snr_threshold, num_img)
2028
+
2029
+ while n < size:
2030
+ # disable print statements
2031
+ with contextlib.redirect_stdout(None):
2032
+ self.dict_buffer = None # this is used to store the sampled lensed_param in batches when running the sampling_routine
2033
+ lensed_param = self.lensed_sampling_routine(
2034
+ size=self.batch_size, output_jsonfile=buffer_file, resume=False
2035
+ ) # Dimensions are (size, n_max_images)
2036
+
2037
+ total_events_in_this_iteration = len(lensed_param["zs"])
2038
+
2039
+ # Below code ensures that the snr is recalculated for the detectable events with inner product
2040
+ # The code is use when the snr is calculated with 'ann' method of `gwsnr`
2041
+ if snr_recalculation:
2042
+ lensed_param = self._recalculate_snr_lensed(lensed_param, snr_threshold_recalculation, num_img, total_events_in_this_iteration)
2043
+
2044
+ snr_hit = self._find_detectable_index_lensed(lensed_param, snr_threshold, pdet_threshold, num_img, detectability_condition)
2045
+
2046
+ # store all params in json file
2047
+ 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)
2048
+
2049
+ n += np.sum(snr_hit)
2050
+ events_total += total_events_in_this_iteration
2051
+ total_rate = self.rate_function(n, events_total, param_type="lensed", verbose=False)
2052
+
2053
+ # save meta data
2054
+ self._append_meta_data(meta_data_path, n, events_total, total_rate)
2055
+
2056
+ print(f"storing detectable lensed params in {output_path}")
2057
+ print(f"storing meta data in {meta_data_path}")
2058
+
2059
+ if trim_to_size:
2060
+ param_final, total_rate = self._trim_results_to_size(size, output_path, meta_data_path, param_type="lensed")
2061
+ else:
2062
+ param_final = get_param_from_json(output_path)
2063
+
2064
+ # call self.json_file_names["ler_param"] and for adding the final results
2065
+ data = load_json(self.ler_directory+"/"+self.json_file_names["ler_params"])
2066
+ # write the results
2067
+ try:
2068
+ data["detectable_lensed_rate_per_year"] = total_rate
2069
+ data["detectability_condition_lensed"] = detectability_condition
2070
+ except:
2071
+ meta = get_param_from_json(meta_data_path)
2072
+ data["detectable_lensed_rate_per_year"] = meta["total_rate"][-1]
2073
+ data["detectability_condition_lensed"] = detectability_condition
2074
+ append_json(self.ler_directory+"/"+self.json_file_names["ler_params"], data, replace=True)
2075
+
2076
+ return param_final
2077
+
2078
+ def _initial_setup_for_n_event_selection(self, meta_data_file, output_jsonfile, resume, batch_size):
2079
+ """Helper function for selecting_n_unlensed_detectable_events and selecting_n_lensed_detectable_events functions.
2080
+
2081
+ Parameters
2082
+ ----------
2083
+ meta_data_file : `str`
2084
+ json file name for storing the metadata.
2085
+ output_jsonfile : `str`
2086
+ json file name for storing the parameters of the detectable events.
2087
+ resume : `bool`
2088
+ resume = False (default) or True.
2089
+ if True, the function will resume from the last batch.
2090
+ batch_size : `int`
2091
+ batch size for sampling.
2092
+ default batch_size = 50000.
2093
+
2094
+ Returns
2095
+ ----------
2096
+ n : `int`
2097
+ iterator.
2098
+ events_total : `int`
2099
+ total number of events.
2100
+ output_path : `str`
2101
+ path to the output json file.
2102
+ meta_data_path : `str`
2103
+ path to the metadata json file.
2104
+ buffer_file : `str`
2105
+ path to the buffer json file.
1673
2106
  """
1674
2107
 
1675
2108
  meta_data_path = self.ler_directory+"/"+meta_data_file
1676
2109
  output_path = self.ler_directory+"/"+output_jsonfile
1677
2110
  if meta_data_path==output_path:
1678
2111
  raise ValueError("meta_data_file and output_jsonfile cannot be same.")
1679
-
2112
+
1680
2113
  if batch_size is None:
1681
2114
  batch_size = self.batch_size
1682
2115
  else:
@@ -1691,187 +2124,105 @@ class LeR(LensGalaxyParameterDistribution):
1691
2124
  os.remove(meta_data_path)
1692
2125
  else:
1693
2126
  # get sample size as size from json file
1694
- try:
2127
+ if os.path.exists(output_path):
1695
2128
  param_final = get_param_from_json(output_path)
1696
- n = load_json(meta_data_path)["detectable_events"][-1]
2129
+ n = len(param_final["zs"])
1697
2130
  events_total = load_json(meta_data_path)["events_total"][-1]
1698
2131
  del param_final
1699
- except:
2132
+ else:
1700
2133
  n = 0
1701
2134
  events_total = 0
1702
2135
 
1703
- # check for images with snr above threshold
1704
- # convert to array
1705
- snr_threshold = np.array([snr_threshold]).reshape(-1)
1706
- num_img = np.array([num_img]).reshape(-1)
1707
- # get descending sorted idx of snr_threshold
1708
- idx = np.argsort(-snr_threshold)
1709
- snr_threshold = snr_threshold[idx]
1710
- num_img = num_img[idx]
1711
-
1712
- buffer_file = self.ler_directory+"/"+"lensed_params_buffer.json"
2136
+ buffer_file = "params_buffer.json"
1713
2137
  print("collected number of detectable events = ", n)
1714
- # loop until n samples are collected
1715
- while n < size:
1716
- # disable print statements
1717
- with contextlib.redirect_stdout(None):
1718
- self.dict_buffer = None
1719
- lensed_param = self.lensed_sampling_routine(
1720
- size=self.batch_size, output_jsonfile=buffer_file, resume=False
1721
- ) # Dimensions are (size, n_max_images)
1722
-
1723
- total_events_in_this_iteration = len(lensed_param["zs"])
1724
- ############################
1725
- # for ANN SNR #
1726
- ############################
1727
- if snr_recalculation:
1728
- # dealing with provided snr_threshold_recalculation
1729
- snr_threshold_recalculation = np.array([snr_threshold_recalculation]).reshape(-1)
1730
- idx = np.argsort(-snr_threshold_recalculation)
1731
- snr_threshold_recalculation = snr_threshold_recalculation[idx]
1732
- num_img_recalculation = num_img[idx]
1733
-
1734
- # check optimal_snr_net is provided in dict
1735
- snr_param = lensed_param["optimal_snr_net"]
1736
- if "optimal_snr_net" in lensed_param.keys():
1737
- snr_param = -np.sort(-snr_param, axis=1) # sort snr in descending order
1738
- else:
1739
- raise ValueError("optimal_snr_net not provided in lensed_param dict. Exiting...")
1740
2138
 
1741
- # 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
1742
- j = 0
1743
- idx_max = 0
1744
- snr_hit = np.full(len(snr_param), True) # boolean array to store the result of the threshold condition
1745
- for i, snr_th in enumerate(snr_threshold_recalculation):
1746
- idx_max = idx_max + num_img_recalculation[i]
2139
+ return n, events_total, output_path, meta_data_path, buffer_file
1747
2140
 
1748
- snr_hit = snr_hit & (np.sum((snr_param[:,j:idx_max] > snr_th), axis=1) >= num_img_recalculation[i])
1749
- j = idx_max
2141
+ def _trim_results_to_size(self, size, output_path, meta_data_path, param_type="unlensed"):
2142
+ """
2143
+ Helper function of 'selecting_n_unlensed_detectable_events' and 'selecting_n_lensed_detectable_events' functions. Trims the data in the output file to the specified size and updates the metadata accordingly.
1750
2144
 
1751
- # reduce the size of the dict
1752
- for key, value in lensed_param.items():
1753
- lensed_param[key] = value[snr_hit]
1754
- # recalculate more accurate snrs
1755
- print("calculating snrs...")
1756
- snrs, lensed_param = self.get_lensed_snrs(
1757
- snr_calculator=self.snr_bilby,
1758
- list_of_detectors=self.list_of_detectors,
1759
- lensed_param=lensed_param,
1760
- )
1761
- lensed_param.update(snrs)
1762
- ############################
2145
+ Parameters
2146
+ ----------
2147
+ size : `int`
2148
+ number of samples to be selected.
2149
+ output_path : `str`
2150
+ path to the output json file.
2151
+ meta_data_path : `str`
2152
+ path to the metadata json file.
2153
+ param_type : `str`
2154
+ type of parameters.
2155
+ default param_type = "unlensed".
2156
+ other options are "lensed".
1763
2157
 
1764
- if detectability_condition == "step_function":
1765
- snr_hit = np.full(len(lensed_param["zs"]), True) # boolean array to store the result of the threshold condition
1766
- try:
1767
- snr_param = lensed_param["optimal_snr_net"]
1768
- snr_param = -np.sort(-snr_param, axis=1) # sort snr in descending order
1769
- except:
1770
- raise ValueError("snr not provided in lensed_param dict. Exiting...")
1771
-
1772
- # 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
1773
- j = 0
1774
- idx_max = 0
1775
- for i in range(len(snr_threshold)):
1776
- idx_max = idx_max + num_img[i]
1777
- snr_hit = snr_hit & (np.sum((snr_param[:,j:idx_max] > snr_threshold[i]), axis=1) >= num_img[i])
1778
- j = idx_max
2158
+ Returns
2159
+ ----------
2160
+ param_final : `dict`
2161
+ dictionary of GW source parameters of the detectable events. Refer to :attr:`~unlensed_param` or :attr:`~lensed_param` for details.
2162
+ new_total_rate : `float`
2163
+ total rate (Mpc^-3 yr^-1).
2164
+ """
1779
2165
 
1780
- elif detectability_condition == "pdet":
1781
- # check if pdet is provided in unlensed_param dict
1782
- if "pdet_net" in lensed_param.keys():
1783
- pdet = lensed_param["pdet_net"]
1784
- pdet = -np.sort(-pdet, axis=1) # sort pdet in descending order
1785
- pdet = pdet[:,:np.sum(num_img)] # use only num_img images
1786
- else:
1787
- if "optimal_snr_net" in lensed_param.keys():
1788
- # pdet dimension is (size, n_max_images)
1789
- snr_param = lensed_param["optimal_snr_net"]
1790
- snr_param = -np.sort(-snr_param, axis=1) # sort snr in descending order
1791
-
1792
- pdet = np.ones(np.shape(snr_param))
1793
- j = 0
1794
- idx_max = 0
1795
- for i in range(len(snr_threshold)):
1796
- idx_max = idx_max + num_img[i]
1797
- pdet[:,j:idx_max] = 1 - norm.cdf(snr_threshold[i] - snr_param[:,j:idx_max])
1798
- j = idx_max
1799
- else:
1800
- print("pdet or optimal_snr_net not provided in lensed_param dict. Exiting...")
1801
- return None
1802
-
1803
- snr_hit = np.prod(pdet, axis=1)>0.5
1804
-
1805
- # store all params in json file
1806
- if nan_to_num:
1807
- for key, value in lensed_param.items():
1808
- lensed_param[key] = np.nan_to_num(value[snr_hit])
1809
- else:
1810
- for key, value in lensed_param.items():
1811
- lensed_param[key] = value[snr_hit]
2166
+ print(f"\n trmming final result to size={size}")
2167
+ param_final = get_param_from_json(output_path)
2168
+ # randomly select size number of samples
2169
+ len_ = len(list(param_final.values())[0])
2170
+ idx = np.random.choice(len_, size, replace=False)
2171
+ # trim the final param dictionary, randomly, without repeating
2172
+ for key, value in param_final.items():
2173
+ param_final[key] = value[idx]
2174
+
2175
+ # change meta data
2176
+ meta_data = load_json(meta_data_path)
2177
+ old_events_total = meta_data["events_total"][-1]
2178
+ old_detectable_events = meta_data["detectable_events"][-1]
2179
+
2180
+ # adjust the meta data
2181
+ # following is to keep rate the same
2182
+ new_events_total = np.round(size*old_events_total/old_detectable_events)
2183
+ new_total_rate = self.rate_function(size, new_events_total, param_type=param_type, verbose=False)
2184
+ meta_data["events_total"][-1] = new_events_total
2185
+ meta_data["detectable_events"][-1] = size
2186
+ meta_data["total_rate"][-1] = new_total_rate
2187
+
2188
+ print("collected number of detectable events = ", size)
2189
+ print("total number of events = ", new_events_total)
2190
+ print(f"total {param_type} event rate (yr^-1): {new_total_rate}")
2191
+
2192
+ # save the meta data
2193
+ append_json(meta_data_path, meta_data, replace=True)
2194
+ # save the final param dictionary
2195
+ append_json(output_path, param_final, replace=True)
2196
+
2197
+ return param_final, new_total_rate
2198
+
2199
+ def _append_meta_data(self, meta_data_path, n, events_total, total_rate):
2200
+ """
2201
+ Helper function for appending meta data json file.
1812
2202
 
1813
- if os.path.exists(output_path):
1814
- try:
1815
- append_json(output_path, lensed_param, replace=False)
1816
- except:
1817
- append_json(output_path, lensed_param, replace=True)
1818
- else:
1819
- append_json(output_path, lensed_param, replace=True)
2203
+ Parameters
2204
+ ----------
2205
+ meta_data_path : `str`
2206
+ path to the metadata json file.
2207
+ n : `int`
2208
+ iterator.
2209
+ events_total : `int`
2210
+ total number of events.
2211
+ total_rate : `float`
2212
+ total rate (Mpc^-3 yr^-1).
2213
+ """
1820
2214
 
1821
- n += np.sum(snr_hit)
1822
- events_total += total_events_in_this_iteration
2215
+ # save meta data
2216
+ meta_data = dict(events_total=[events_total], detectable_events=[float(n)], total_rate=[total_rate])
1823
2217
 
1824
- # save meta data
1825
- total_rate = self.normalization_pdf_z_lensed * n / events_total
1826
- meta_data = dict(events_total=[events_total], detectable_events=[float(n)], total_rate=[total_rate])
1827
- if os.path.exists(meta_data_path):
2218
+ if os.path.exists(meta_data_path):
2219
+ try:
1828
2220
  append_json(meta_data_path, meta_data, replace=False)
1829
- else:
2221
+ except:
1830
2222
  append_json(meta_data_path, meta_data, replace=True)
1831
-
1832
- print("collected number of detectable events = ", n)
1833
- print("total number of events = ", events_total)
1834
- print(f"total lensed rate (yr^-1): {total_rate}")
1835
-
1836
- print(f"storing detectable lensed params in {output_path}")
1837
- print(f"storing meta data in {meta_data_path}")
1838
-
1839
- if trim_to_size:
1840
- # trim the final param dictionary
1841
- print(f"\n trmming final result to size={size}")
1842
- param_final = get_param_from_json(output_path)
1843
- # randomly select size number of samples
1844
- len_ = len(list(param_final.values())[0])
1845
- idx = np.random.choice(len_, size, replace=False)
1846
- # trim the final param dictionary
1847
- for key, value in param_final.items():
1848
- param_final[key] = value[idx]
1849
-
1850
- # change meta data
1851
- meta_data = load_json(meta_data_path)
1852
- old_events_total = meta_data["events_total"][-1]
1853
- old_detectable_events = meta_data["detectable_events"][-1]
1854
-
1855
- # adjust the meta data
1856
- # following is to keep rate the same
1857
- new_events_total = np.round(size * old_events_total/old_detectable_events)
1858
- new_total_rate = self.normalization_pdf_z_lensed * size / new_events_total
1859
- meta_data["events_total"][-1] = new_events_total
1860
- meta_data["detectable_events"][-1] = size
1861
- meta_data["total_rate"][-1] = new_total_rate
1862
-
1863
- print("collected number of detectable events = ", size)
1864
- print("total number of events = ", new_events_total)
1865
- print(f"total unlensed rate (yr^-1): {new_total_rate}")
1866
-
1867
- # save the meta data
1868
- append_json(meta_data_path, meta_data, replace=True)
1869
- # save the final param dictionary
1870
- append_json(output_path, param_final, replace=True)
1871
2223
  else:
1872
- param_final = get_param_from_json(output_path)
1873
-
1874
- return param_final
1875
-
1876
-
2224
+ append_json(meta_data_path, meta_data, replace=True)
1877
2225
 
2226
+ print("collected number of detectable events = ", n)
2227
+ print("total number of events = ", events_total)
2228
+ print(f"total rate (yr^-1): {total_rate}")